diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index d37712c682c2..e9a13fdac2a8 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -421,9 +421,12 @@ public static class QueryPaginator { QueryPaginator(@Nonnull Query query, int pageSize) { this.hasOverallLimit = query.builder.getRowsLimit() > 0; this.remainingRows = query.builder.getRowsLimit(); - this.query = query.limit(pageSize); + this.query = query.clone(); if (hasOverallLimit) { + this.query.limit(Math.min(pageSize, remainingRows)); remainingRows -= pageSize; + } else { + this.query.limit(pageSize); } this.pageSize = pageSize; this.prevSplitPoint = ByteString.EMPTY; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index 1ed9ce1c59dd..b63be7ae3296 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -635,6 +635,7 @@ public void testQueryPaginatorRangeLimitReached() { // Table query test. int chunkSize = 10, limit = 15; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); + Query originalCopy = query.clone(); Query.QueryPaginator paginator = query.createPaginator(chunkSize); Query nextQuery = paginator.getNextQuery(); @@ -667,6 +668,7 @@ public void testQueryPaginatorRangeLimitReached() { assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); assertThat(paginator.advance(ByteString.copyFromUtf8("d"))).isFalse(); + assertThat(query).isEqualTo(originalCopy); // AuthorizedView query test. query = Query.create(AUTHORIZED_VIEW_ID).range("a", "z").limit(limit); @@ -967,6 +969,31 @@ public void testQueryPaginatorEmptyTable() { assertThat(queryPaginator.advance(ByteString.EMPTY)).isFalse(); } + @Test + public void testQueryPaginatorLimitSmallerThanChunkSize() { + int chunkSize = 100, limit = 5; + Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); + Query originalCopy = query.clone(); + + Query.QueryPaginator paginator = query.createPaginator(chunkSize); + Query nextQuery = paginator.getNextQuery(); + + ReadRowsRequest expectedProto = + expectedReadFromTableProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("a")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(limit) + .build(); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto); + assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isFalse(); + assertThat(query).isEqualTo(originalCopy); + } + @Test public void testQueryReversed() { Query query = Query.create(TABLE_ID).reversed(true);