Skip to content

Commit 413031d

Browse files
mutianfmutianf
committed
feat(bigtable): route point read rows to shim
1 parent 8cf6b07 commit 413031d

5 files changed

Lines changed: 435 additions & 12 deletions

File tree

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -318,10 +318,39 @@ public TargetId getTargetId() {
318318
return targetId;
319319
}
320320

321+
/**
322+
* Returns true if this query identifies a single row that can be served by a point read. Supports
323+
* two shapes: exactly one row key and no row ranges, or exactly one closed-closed row range whose
324+
* start key equals its end key.
325+
*/
326+
@InternalApi
327+
public boolean isSinglePointQuery() {
328+
RowSet rows = this.builder.getRows();
329+
int keyCount = rows.getRowKeysCount();
330+
int rangeCount = rows.getRowRangesCount();
331+
if (keyCount == 1 && rangeCount == 0) {
332+
return true;
333+
}
334+
if (keyCount == 0 && rangeCount == 1) {
335+
RowRange range = rows.getRowRanges(0);
336+
return range.hasStartKeyClosed()
337+
&& range.hasEndKeyClosed()
338+
&& range.getStartKeyClosed().equals(range.getEndKeyClosed());
339+
}
340+
return false;
341+
}
342+
321343
@InternalApi
322344
public SessionReadRowRequest toSessionPointProto() {
345+
Preconditions.checkState(
346+
isSinglePointQuery(),
347+
"Query must be a single-point read (one row key, or one closed-closed row range whose"
348+
+ " start equals its end)");
349+
RowSet rows = this.builder.getRows();
350+
ByteString key =
351+
rows.getRowKeysCount() > 0 ? rows.getRowKeys(0) : rows.getRowRanges(0).getStartKeyClosed();
323352
return SessionReadRowRequest.newBuilder()
324-
.setKey(this.builder.getRows().getRowKeysList().get(0))
353+
.setKey(key)
325354
.setFilter(this.builder.getFilter())
326355
.build();
327356
}

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

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
2828
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
2929
import com.google.api.gax.retrying.RetryAlgorithm;
30+
import com.google.api.gax.retrying.RetrySettings;
3031
import com.google.api.gax.retrying.RetryingExecutorWithContext;
3132
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
3233
import com.google.api.gax.retrying.SimpleStreamResumptionStrategy;
@@ -97,6 +98,7 @@
9798
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
9899
import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
99100
import com.google.cloud.bigtable.data.v2.stub.readrows.LargeReadRowsResumptionStrategy;
101+
import com.google.cloud.bigtable.data.v2.stub.readrows.MaybePointReadCallable;
100102
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
101103
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
102104
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
@@ -259,11 +261,17 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
259261
bigtableClientContext.getClientContext().getTracerFactory(),
260262
span);
261263

262-
return traced.withDefaultCallContext(
263-
bigtableClientContext
264-
.getClientContext()
265-
.getDefaultCallContext()
266-
.withRetrySettings(perOpSettings.readRowsSettings.getRetrySettings()));
264+
ServerStreamingCallable<Query, RowT> classic =
265+
traced.withDefaultCallContext(
266+
bigtableClientContext
267+
.getClientContext()
268+
.getDefaultCallContext()
269+
.withRetrySettings(perOpSettings.readRowsSettings.getRetrySettings()));
270+
271+
return new MaybePointReadCallable<>(
272+
classic,
273+
createPointReadCallable(
274+
rowAdapter, "ReadRows", perOpSettings.readRowsSettings.getRetrySettings()));
267275
}
268276

269277
/**
@@ -281,13 +289,19 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
281289
* </ul>
282290
*/
283291
public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT> rowAdapter) {
292+
return createPointReadCallable(
293+
rowAdapter, "ReadRow", perOpSettings.readRowSettings.getRetrySettings());
294+
}
295+
296+
private <RowT> UnaryCallable<Query, RowT> createPointReadCallable(
297+
RowAdapter<RowT> rowAdapter, String spanName, RetrySettings retrySettings) {
284298
ClientContext clientContext = bigtableClientContext.getClientContext();
285299

286300
ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
287301
createReadRowsBaseCallable(
288302
ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
289303
.setRetryableCodes(perOpSettings.readRowSettings.getRetryableCodes())
290-
.setRetrySettings(perOpSettings.readRowSettings.getRetrySettings())
304+
.setRetrySettings(retrySettings)
291305
.setIdleTimeoutDuration(Duration.ZERO)
292306
.setWaitTimeoutDuration(Duration.ZERO)
293307
.build(),
@@ -302,16 +316,17 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
302316
BigtableUnaryOperationCallable<Query, RowT> classic =
303317
new BigtableUnaryOperationCallable<>(
304318
readRowCallable,
305-
clientContext
306-
.getDefaultCallContext()
307-
.withRetrySettings(perOpSettings.readRowSettings.getRetrySettings()),
319+
clientContext.getDefaultCallContext().withRetrySettings(retrySettings),
308320
clientContext.getTracerFactory(),
309-
getSpanName("ReadRow"),
321+
getSpanName(spanName),
310322
/* allowNoResponse= */ true);
311323

324+
UnaryCallSettings<?, ?> shimSettings =
325+
perOpSettings.readRowSettings.toBuilder().setRetrySettings(retrySettings).build();
326+
312327
return bigtableClientContext
313328
.getSessionShim()
314-
.decorateReadRow(classic, rowAdapter, perOpSettings.readRowSettings);
329+
.decorateReadRow(classic, rowAdapter, shimSettings);
315330
}
316331

317332
private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRowsBaseCallable(

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

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub.readrows;
17+
18+
import com.google.api.core.ApiFuture;
19+
import com.google.api.core.ApiFutureCallback;
20+
import com.google.api.core.ApiFutures;
21+
import com.google.api.core.InternalApi;
22+
import com.google.api.gax.rpc.ApiCallContext;
23+
import com.google.api.gax.rpc.ResponseObserver;
24+
import com.google.api.gax.rpc.ServerStreamingCallable;
25+
import com.google.api.gax.rpc.StreamController;
26+
import com.google.api.gax.rpc.UnaryCallable;
27+
import com.google.cloud.bigtable.data.v2.models.Query;
28+
import com.google.common.util.concurrent.MoreExecutors;
29+
30+
/**
31+
* Routes ReadRows calls whose query identifies a single row through a unary point-read callable,
32+
* letting them benefit from the same session-shim diversion as {@code ReadRow}. Queries that cannot
33+
* be reduced to a point read fall through to the classic {@code ReadRows} callable.
34+
*/
35+
@InternalApi
36+
public class MaybePointReadCallable<RowT> extends ServerStreamingCallable<Query, RowT> {
37+
private final ServerStreamingCallable<Query, RowT> classic;
38+
private final UnaryCallable<Query, RowT> pointReader;
39+
40+
public MaybePointReadCallable(
41+
ServerStreamingCallable<Query, RowT> classic, UnaryCallable<Query, RowT> pointReader) {
42+
this.classic = classic;
43+
this.pointReader = pointReader;
44+
}
45+
46+
@Override
47+
public void call(Query request, ResponseObserver<RowT> responseObserver, ApiCallContext context) {
48+
if (!request.isSinglePointQuery()) {
49+
classic.call(request, responseObserver, context);
50+
return;
51+
}
52+
53+
ApiFuture<RowT> future = pointReader.futureCall(request, context);
54+
responseObserver.onStart(
55+
new StreamController() {
56+
@Override
57+
public void cancel() {
58+
future.cancel(false);
59+
}
60+
61+
@Override
62+
public void disableAutoInboundFlowControl() {}
63+
64+
@Override
65+
public void request(int count) {}
66+
});
67+
68+
ApiFutures.addCallback(
69+
future,
70+
new ApiFutureCallback<RowT>() {
71+
@Override
72+
public void onSuccess(RowT row) {
73+
if (row != null) {
74+
try {
75+
responseObserver.onResponse(row);
76+
} catch (Throwable t) {
77+
responseObserver.onError(t);
78+
return;
79+
}
80+
}
81+
responseObserver.onComplete();
82+
}
83+
84+
@Override
85+
public void onFailure(Throwable t) {
86+
responseObserver.onError(t);
87+
}
88+
},
89+
MoreExecutors.directExecutor());
90+
}
91+
}

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

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.bigtable.v2.RowFilter;
2424
import com.google.bigtable.v2.RowRange;
2525
import com.google.bigtable.v2.RowSet;
26+
import com.google.bigtable.v2.SessionReadRowRequest;
2627
import com.google.cloud.bigtable.data.v2.internal.ByteStringComparator;
2728
import com.google.cloud.bigtable.data.v2.internal.NameUtil;
2829
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
@@ -966,4 +967,107 @@ public void testQueryReversed() {
966967
assertThat(query.toProto(requestContext))
967968
.isEqualTo(expectedReadFromTableProtoBuilder().setReversed(true).build());
968969
}
970+
971+
@Test
972+
public void isSinglePointQuery_singleRowKey() {
973+
assertThat(Query.create(TABLE_ID).rowKey("k").isSinglePointQuery()).isTrue();
974+
}
975+
976+
@Test
977+
public void isSinglePointQuery_singleClosedRange() {
978+
assertThat(
979+
Query.create(TABLE_ID)
980+
.range(ByteStringRange.unbounded().startClosed("k").endClosed("k"))
981+
.isSinglePointQuery())
982+
.isTrue();
983+
}
984+
985+
@Test
986+
public void isSinglePointQuery_emptyQuery() {
987+
assertThat(Query.create(TABLE_ID).isSinglePointQuery()).isFalse();
988+
}
989+
990+
@Test
991+
public void isSinglePointQuery_multipleRowKeys() {
992+
assertThat(Query.create(TABLE_ID).rowKey("a").rowKey("b").isSinglePointQuery()).isFalse();
993+
}
994+
995+
@Test
996+
public void isSinglePointQuery_rowKeyAndRange() {
997+
assertThat(
998+
Query.create(TABLE_ID)
999+
.rowKey("a")
1000+
.range(ByteStringRange.unbounded().startClosed("a").endClosed("a"))
1001+
.isSinglePointQuery())
1002+
.isFalse();
1003+
}
1004+
1005+
@Test
1006+
public void isSinglePointQuery_multipleRanges() {
1007+
assertThat(
1008+
Query.create(TABLE_ID)
1009+
.range(ByteStringRange.unbounded().startClosed("a").endClosed("a"))
1010+
.range(ByteStringRange.unbounded().startClosed("b").endClosed("b"))
1011+
.isSinglePointQuery())
1012+
.isFalse();
1013+
}
1014+
1015+
@Test
1016+
public void isSinglePointQuery_closedOpenRange() {
1017+
assertThat(
1018+
Query.create(TABLE_ID)
1019+
.range(ByteStringRange.unbounded().startClosed("k").endOpen("k"))
1020+
.isSinglePointQuery())
1021+
.isFalse();
1022+
}
1023+
1024+
@Test
1025+
public void isSinglePointQuery_unequalClosedRange() {
1026+
assertThat(
1027+
Query.create(TABLE_ID)
1028+
.range(ByteStringRange.unbounded().startClosed("a").endClosed("b"))
1029+
.isSinglePointQuery())
1030+
.isFalse();
1031+
}
1032+
1033+
@Test
1034+
public void isSinglePointQuery_prefixRange() {
1035+
assertThat(Query.create(TABLE_ID).prefix("k").isSinglePointQuery()).isFalse();
1036+
}
1037+
1038+
@Test
1039+
public void toSessionPointProto_fromRowKey() {
1040+
Query query = Query.create(TABLE_ID).rowKey("the-key");
1041+
assertThat(query.toSessionPointProto())
1042+
.isEqualTo(
1043+
SessionReadRowRequest.newBuilder().setKey(ByteString.copyFromUtf8("the-key")).build());
1044+
}
1045+
1046+
@Test
1047+
public void toSessionPointProto_fromClosedRange() {
1048+
Query query =
1049+
Query.create(TABLE_ID)
1050+
.range(ByteStringRange.unbounded().startClosed("the-key").endClosed("the-key"));
1051+
assertThat(query.toSessionPointProto())
1052+
.isEqualTo(
1053+
SessionReadRowRequest.newBuilder().setKey(ByteString.copyFromUtf8("the-key")).build());
1054+
}
1055+
1056+
@Test
1057+
public void toSessionPointProto_preservesFilter() {
1058+
RowFilter filter = FILTERS.key().regex("regex").toProto();
1059+
Query query = Query.create(TABLE_ID).rowKey("the-key").filter(FILTERS.key().regex("regex"));
1060+
assertThat(query.toSessionPointProto())
1061+
.isEqualTo(
1062+
SessionReadRowRequest.newBuilder()
1063+
.setKey(ByteString.copyFromUtf8("the-key"))
1064+
.setFilter(filter)
1065+
.build());
1066+
}
1067+
1068+
@Test
1069+
public void toSessionPointProto_rejectsNonSinglePointQuery() {
1070+
Query query = Query.create(TABLE_ID).rowKey("a").rowKey("b");
1071+
assertThrows(IllegalStateException.class, query::toSessionPointProto);
1072+
}
9691073
}

0 commit comments

Comments
 (0)