From a59bbb834752d7f42fb97ee427bd6caf9fc62db7 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 22 Nov 2022 12:49:26 -0500 Subject: [PATCH 1/9] feat: add a query paginator --- .../cloud/bigtable/data/v2/models/Query.java | 69 +++++++++++ .../bigtable/data/v2/models/QueryTest.java | 116 ++++++++++++++++++ 2 files changed, 185 insertions(+) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index 986a0ca1a5..b00963d1b0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -15,6 +15,7 @@ */ package com.google.cloud.bigtable.data.v2.models; +import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.bigtable.v2.ReadRowsRequest; import com.google.bigtable.v2.RowFilter; @@ -38,6 +39,7 @@ import java.util.List; import java.util.SortedSet; import javax.annotation.Nonnull; +import javax.annotation.Nullable; /** A simple wrapper to construct a query for the ReadRows RPC. */ public final class Query implements Serializable { @@ -248,6 +250,29 @@ public List shard(SortedSet splitPoints) { return shards; } + /** + * Create a query paginator that'll split the query into smaller chunks. + * + *

Example usage: + * + * 

{@code
+   * Query query = Query.create(...).range("a", "z");
+   * Query.QueryPaginator paginator = query.createQueryPaginator(100);
+   * ByteString lastSeenRowKey = null;
+   * while (paginator.advance(lastSeenRowKey)) {
+   *     List rows = client.readRowsCallable().all().call(paginator.getNextQuery());
+   *     for (Row row : rows) {
+   *        // do some processing
+   *        lastSeenRow = row;
+   *     }
+   * }
+   * }
+ */ + @BetaApi("This surface is stable yet it might be removed in the future.") + public QueryPaginator createQueryPaginator(int chunkSize) { + return new QueryPaginator(this, chunkSize); + } + /** Get the minimal range that encloses all of the row keys and ranges in this Query. */ public ByteStringRange getBound() { return RowSetUtil.getBound(builder.getRows()); @@ -297,6 +322,50 @@ private static ByteString wrapKey(String key) { return ByteString.copyFromUtf8(key); } + @BetaApi("This surface is stable yet it might be removed in the future.") + public class QueryPaginator { + + private long originalLimit; + private long newLimit; + private Query query; + private int chunkSize; + + QueryPaginator(@Nonnull Query query, int chunkSize) { + this.originalLimit = query.builder.getRowsLimit(); + this.newLimit = query.builder.getRowsLimit(); + this.query = query; + this.chunkSize = chunkSize; + } + + Query getNextQuery() { + return query; + } + + boolean advance(@Nullable ByteString lastSeenRowKey) { + if (originalLimit != 0 && newLimit <= 0) { + return false; + } + if (originalLimit != 0) { + query.limit(Math.min(this.chunkSize, newLimit)); + newLimit -= chunkSize; + } else { + query.limit(chunkSize); + } + + ByteString splitPoint = ByteString.EMPTY; + if (lastSeenRowKey != null) { + splitPoint = lastSeenRowKey; + } + RowSetUtil.Split split = RowSetUtil.split(query.builder.getRows(), splitPoint); + if (split.getRight() == null) { + return false; + } + query.builder.setRows(split.getRight()); + + return true; + } + } + @Override public boolean equals(Object o) { if (this == o) { diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index ccb0441c71..f1307d64c7 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -327,4 +327,120 @@ public void testClone() { assertThat(clonedReq).isEqualTo(query); assertThat(clonedReq.toProto(requestContext)).isEqualTo(request); } + + @Test + public void testQueryPaginatorRangeLimitReached() { + int chunkSize = 10, limit = 15; + Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); + Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + paginator.advance(null); + + Query nextQuery = paginator.getNextQuery(); + + Builder expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("a")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(chunkSize); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + paginator.advance(ByteString.copyFromUtf8("c")); + int expectedLimit = limit - chunkSize; + nextQuery = paginator.getNextQuery(); + expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyOpen(ByteString.copyFromUtf8("c")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(expectedLimit); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + assertThat(paginator.advance(ByteString.copyFromUtf8("d"))).isFalse(); + } + + @Test + public void testQueryPaginatorRangeLimitMultiplyOfChunkSize() { + int chunkSize = 10, limit = 20; + Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); + Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + assertThat(paginator.advance(null)).isTrue(); + + Query nextQuery = paginator.getNextQuery(); + + Builder expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("a")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(chunkSize); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isTrue(); + int expectedLimit = limit - chunkSize; + nextQuery = paginator.getNextQuery(); + expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyOpen(ByteString.copyFromUtf8("c")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(expectedLimit); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + assertThat(paginator.advance(ByteString.copyFromUtf8("d"))).isFalse(); + } + + @Test + public void testQueryPaginatorRagneNoLimit() { + int chunkSize = 10; + Query query = Query.create(TABLE_ID).range("a", "z"); + Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + paginator.advance(null); + + Query nextQuery = paginator.getNextQuery(); + + Builder expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyClosed(ByteString.copyFromUtf8("a")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(chunkSize); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + paginator.advance(ByteString.copyFromUtf8("c")); + nextQuery = paginator.getNextQuery(); + expectedProto = + expectedProtoBuilder() + .setRows( + RowSet.newBuilder() + .addRowRanges( + RowRange.newBuilder() + .setStartKeyOpen(ByteString.copyFromUtf8("c")) + .setEndKeyOpen(ByteString.copyFromUtf8("z")) + .build())) + .setRowsLimit(chunkSize); + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + assertThat(paginator.advance(ByteString.copyFromUtf8("z"))).isFalse(); + } } From c1820899fd90374a7c0187e7c7c5c4369378d468 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 22 Nov 2022 20:29:16 -0500 Subject: [PATCH 2/9] add some comments --- .../cloud/bigtable/data/v2/models/Query.java | 9 +++++ .../bigtable/data/v2/models/QueryTest.java | 39 +++++++++++++++++-- 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index b00963d1b0..c59a4b580a 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -322,6 +322,10 @@ private static ByteString wrapKey(String key) { return ByteString.copyFromUtf8(key); } + /** + * A Query Paginator that will split a query into small chunks. See {@link + * Query#createQueryPaginator(int)} for example usage. + */ @BetaApi("This surface is stable yet it might be removed in the future.") public class QueryPaginator { @@ -337,10 +341,15 @@ public class QueryPaginator { this.chunkSize = chunkSize; } + /** Return the next query. Needs to be called after advance(). */ Query getNextQuery() { return query; } + /** + * Construct the next query. Return true if there are more queries to return. False if we've + * read everything. + */ boolean advance(@Nullable ByteString lastSeenRowKey) { if (originalLimit != 0 && newLimit <= 0) { return false; diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index f1307d64c7..f763526c17 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -333,7 +333,7 @@ public void testQueryPaginatorRangeLimitReached() { int chunkSize = 10, limit = 15; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); - paginator.advance(null); + assertThat(paginator.advance(null)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -349,7 +349,7 @@ public void testQueryPaginatorRangeLimitReached() { .setRowsLimit(chunkSize); assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); - paginator.advance(ByteString.copyFromUtf8("c")); + assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isTrue(); int expectedLimit = limit - chunkSize; nextQuery = paginator.getNextQuery(); expectedProto = @@ -411,7 +411,7 @@ public void testQueryPaginatorRagneNoLimit() { int chunkSize = 10; Query query = Query.create(TABLE_ID).range("a", "z"); Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); - paginator.advance(null); + assertThat(paginator.advance(null)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -427,7 +427,7 @@ public void testQueryPaginatorRagneNoLimit() { .setRowsLimit(chunkSize); assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); - paginator.advance(ByteString.copyFromUtf8("c")); + assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isTrue(); nextQuery = paginator.getNextQuery(); expectedProto = expectedProtoBuilder() @@ -443,4 +443,35 @@ public void testQueryPaginatorRagneNoLimit() { assertThat(paginator.advance(ByteString.copyFromUtf8("z"))).isFalse(); } + + @Test + public void testQueryPaginatorRowsNoLimit() { + int chunkSize = 10; + Query query = Query.create(TABLE_ID).rowKey("a").rowKey("b").rowKey("c"); + + Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + assertThat(paginator.advance(null)).isTrue(); + + Query nextQuery = paginator.getNextQuery(); + + ReadRowsRequest.Builder expectedProto = expectedProtoBuilder(); + expectedProto + .getRowsBuilder() + .addRowKeys(ByteString.copyFromUtf8("a")) + .addRowKeys(ByteString.copyFromUtf8("b")) + .addRowKeys(ByteString.copyFromUtf8("c")); + expectedProto.setRowsLimit(chunkSize); + + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + paginator.advance(ByteString.copyFromUtf8("b")); + nextQuery = paginator.getNextQuery(); + expectedProto = expectedProtoBuilder(); + expectedProto.getRowsBuilder().addRowKeys(ByteString.copyFromUtf8("c")); + expectedProto.setRowsLimit(chunkSize); + + assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build()); + + assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isFalse(); + } } From 60a73adde96450c1efc736e82ad87e0e1bbf5d9d Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 22 Nov 2022 21:32:41 -0500 Subject: [PATCH 3/9] add a test for full table scan --- .../cloud/bigtable/data/v2/models/Query.java | 37 +++++++--- .../bigtable/data/v2/models/QueryTest.java | 67 +++++++++++++++---- 2 files changed, 79 insertions(+), 25 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index c59a4b580a..f35e0df241 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -39,7 +39,6 @@ import java.util.List; import java.util.SortedSet; import javax.annotation.Nonnull; -import javax.annotation.Nullable; /** A simple wrapper to construct a query for the ReadRows RPC. */ public final class Query implements Serializable { @@ -258,7 +257,7 @@ public List shard(SortedSet splitPoints) { * 
{@code
    * Query query = Query.create(...).range("a", "z");
    * Query.QueryPaginator paginator = query.createQueryPaginator(100);
-   * ByteString lastSeenRowKey = null;
+   * ByteString lastSeenRowKey = ByteString.EMPTY;
    * while (paginator.advance(lastSeenRowKey)) {
    *     List rows = client.readRowsCallable().all().call(paginator.getNextQuery());
    *     for (Row row : rows) {
@@ -329,16 +328,20 @@ private static ByteString wrapKey(String key) {
   @BetaApi("This surface is stable yet it might be removed in the future.")
   public class QueryPaginator {
 
-    private long originalLimit;
+    final private long originalLimit;
     private long newLimit;
     private Query query;
-    private int chunkSize;
+    final private int chunkSize;
+    private ByteString prevSplitPoint;
+    private boolean firstRun;
 
     QueryPaginator(@Nonnull Query query, int chunkSize) {
       this.originalLimit = query.builder.getRowsLimit();
       this.newLimit = query.builder.getRowsLimit();
       this.query = query;
       this.chunkSize = chunkSize;
+      this.prevSplitPoint = ByteString.EMPTY;
+      this.firstRun = true;
     }
 
     /** Return the next query. Needs to be called after advance(). */
@@ -350,7 +353,22 @@ Query getNextQuery() {
      * Construct the next query. Return true if there are more queries to return. False if we've
      * read everything.
      */
-    boolean advance(@Nullable ByteString lastSeenRowKey) {
+    boolean advance(@Nonnull ByteString lastSeenRowKey) {
+      Preconditions.checkArgument(
+          lastSeenRowKey != null, "lastSeenRowKey cannot be null, use ByteString.EMPTY instead.");
+      // Full table scans don't have ranges or limits. Keep track of the previous split
+      // point. If it's the same as the current input return false. The only exception
+      // is the first run for this paginator. So keep track of the first run state too.
+      if (!firstRun && prevSplitPoint.equals(lastSeenRowKey)) {
+        return false;
+      }
+      if (firstRun) {
+        firstRun = false;
+      }
+      this.prevSplitPoint = lastSeenRowKey;
+
+      // Set the query limit. If the original limit is set, return false if the new
+      // limit is <= 0 to avoid returning more rows than intended.
       if (originalLimit != 0 && newLimit <= 0) {
         return false;
       }
@@ -361,16 +379,13 @@ boolean advance(@Nullable ByteString lastSeenRowKey) {
         query.limit(chunkSize);
       }
 
-      ByteString splitPoint = ByteString.EMPTY;
-      if (lastSeenRowKey != null) {
-        splitPoint = lastSeenRowKey;
-      }
-      RowSetUtil.Split split = RowSetUtil.split(query.builder.getRows(), splitPoint);
+      // Split the row ranges / row keys. Return false if there's nothing
+      // left on the right of the split point.
+      RowSetUtil.Split split = RowSetUtil.split(query.builder.getRows(), lastSeenRowKey);
       if (split.getRight() == null) {
         return false;
       }
       query.builder.setRows(split.getRight());
-
       return true;
     }
   }
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
index f763526c17..79fe4a32fe 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java
@@ -333,7 +333,7 @@ public void testQueryPaginatorRangeLimitReached() {
     int chunkSize = 10, limit = 15;
     Query query = Query.create(TABLE_ID).range("a", "z").limit(limit);
     Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize);
-    assertThat(paginator.advance(null)).isTrue();
+    assertThat(paginator.advance(ByteString.EMPTY)).isTrue();
 
     Query nextQuery = paginator.getNextQuery();
 
@@ -372,7 +372,7 @@ public void testQueryPaginatorRangeLimitMultiplyOfChunkSize() {
     int chunkSize = 10, limit = 20;
     Query query = Query.create(TABLE_ID).range("a", "z").limit(limit);
     Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize);
-    assertThat(paginator.advance(null)).isTrue();
+    assertThat(paginator.advance(ByteString.EMPTY)).isTrue();
 
     Query nextQuery = paginator.getNextQuery();
 
@@ -411,7 +411,7 @@ public void testQueryPaginatorRagneNoLimit() {
     int chunkSize = 10;
     Query query = Query.create(TABLE_ID).range("a", "z");
     Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize);
-    assertThat(paginator.advance(null)).isTrue();
+    assertThat(paginator.advance(ByteString.EMPTY)).isTrue();
 
     Query nextQuery = paginator.getNextQuery();
 
@@ -429,16 +429,15 @@ public void testQueryPaginatorRagneNoLimit() {
 
     assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isTrue();
     nextQuery = paginator.getNextQuery();
-    expectedProto =
-        expectedProtoBuilder()
-            .setRows(
-                RowSet.newBuilder()
-                    .addRowRanges(
-                        RowRange.newBuilder()
-                            .setStartKeyOpen(ByteString.copyFromUtf8("c"))
-                            .setEndKeyOpen(ByteString.copyFromUtf8("z"))
-                            .build()))
-            .setRowsLimit(chunkSize);
+    expectedProto
+        .setRows(
+            RowSet.newBuilder()
+                .addRowRanges(
+                    RowRange.newBuilder()
+                        .setStartKeyOpen(ByteString.copyFromUtf8("c"))
+                        .setEndKeyOpen(ByteString.copyFromUtf8("z"))
+                        .build()))
+        .setRowsLimit(chunkSize);
     assertThat(nextQuery.toProto(requestContext)).isEqualTo(expectedProto.build());
 
     assertThat(paginator.advance(ByteString.copyFromUtf8("z"))).isFalse();
@@ -450,7 +449,7 @@ public void testQueryPaginatorRowsNoLimit() {
     Query query = Query.create(TABLE_ID).rowKey("a").rowKey("b").rowKey("c");
 
     Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize);
-    assertThat(paginator.advance(null)).isTrue();
+    assertThat(paginator.advance(ByteString.EMPTY)).isTrue();
 
     Query nextQuery = paginator.getNextQuery();
 
@@ -474,4 +473,44 @@ public void testQueryPaginatorRowsNoLimit() {
 
     assertThat(paginator.advance(ByteString.copyFromUtf8("c"))).isFalse();
   }
+
+  @Test
+  public void testQueryPaginatorFullTableScan() {
+    int chunkSize = 10;
+    Query query = Query.create(TABLE_ID);
+    Query.QueryPaginator queryPaginator = query.createQueryPaginator(chunkSize);
+    assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue();
+
+    ReadRowsRequest.Builder expectedProto =
+        expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize);
+    assertThat(queryPaginator.getNextQuery().toProto(requestContext))
+        .isEqualTo(expectedProto.build());
+
+    assertThat(queryPaginator.advance(ByteString.copyFromUtf8("a"))).isTrue();
+    expectedProto
+        .setRows(
+            RowSet.newBuilder()
+                .addRowRanges(
+                    RowRange.newBuilder().setStartKeyOpen(ByteString.copyFromUtf8("a")).build()))
+        .setRowsLimit(chunkSize);
+    assertThat(queryPaginator.getNextQuery().toProto(requestContext))
+        .isEqualTo(expectedProto.build());
+
+    assertThat(queryPaginator.advance(ByteString.copyFromUtf8("a"))).isFalse();
+  }
+
+  @Test
+  public void testQueryPaginatorEmptyTable() {
+    int chunkSize = 10;
+    Query query = Query.create(TABLE_ID);
+    Query.QueryPaginator queryPaginator = query.createQueryPaginator(chunkSize);
+    assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue();
+
+    ReadRowsRequest.Builder expectedProto =
+        expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize);
+    assertThat(queryPaginator.getNextQuery().toProto(requestContext))
+        .isEqualTo(expectedProto.build());
+
+    assertThat(queryPaginator.advance(ByteString.EMPTY)).isFalse();
+  }
 }

From 72dde6fb2ae5aa94cb8f8a2c5996d4825cfb1d9a Mon Sep 17 00:00:00 2001
From: Mattie Fu 
Date: Tue, 22 Nov 2022 21:34:56 -0500
Subject: [PATCH 4/9] fix format

---
 .../java/com/google/cloud/bigtable/data/v2/models/Query.java  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
index f35e0df241..6c5f13b673 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
@@ -328,10 +328,10 @@ private static ByteString wrapKey(String key) {
   @BetaApi("This surface is stable yet it might be removed in the future.")
   public class QueryPaginator {
 
-    final private long originalLimit;
+    private final long originalLimit;
     private long newLimit;
     private Query query;
-    final private int chunkSize;
+    private final int chunkSize;
     private ByteString prevSplitPoint;
     private boolean firstRun;
 

From 5b6dd8e6448c8e5a5a95da47b6337444647653ca Mon Sep 17 00:00:00 2001
From: Mattie Fu 
Date: Tue, 29 Nov 2022 12:41:40 -0500
Subject: [PATCH 5/9] address comments

---
 .../cloud/bigtable/data/v2/models/Query.java  | 52 ++++++++-----------
 .../bigtable/data/v2/models/QueryTest.java    | 12 ++---
 2 files changed, 29 insertions(+), 35 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
index 6c5f13b673..6d9880c3e4 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java
@@ -268,8 +268,8 @@ public List shard(SortedSet splitPoints) {
    * }
*/ @BetaApi("This surface is stable yet it might be removed in the future.") - public QueryPaginator createQueryPaginator(int chunkSize) { - return new QueryPaginator(this, chunkSize); + public QueryPaginator createPaginator(int pageSize) { + return new QueryPaginator(this, pageSize); } /** Get the minimal range that encloses all of the row keys and ranges in this Query. */ @@ -323,29 +323,27 @@ private static ByteString wrapKey(String key) { /** * A Query Paginator that will split a query into small chunks. See {@link - * Query#createQueryPaginator(int)} for example usage. + * Query#createPaginator(int)} for example usage. */ @BetaApi("This surface is stable yet it might be removed in the future.") - public class QueryPaginator { + public static class QueryPaginator { - private final long originalLimit; - private long newLimit; + private final boolean hasOverallLimit; + private long remainingRows; private Query query; - private final int chunkSize; + private final int pageSize; private ByteString prevSplitPoint; - private boolean firstRun; - QueryPaginator(@Nonnull Query query, int chunkSize) { - this.originalLimit = query.builder.getRowsLimit(); - this.newLimit = query.builder.getRowsLimit(); + QueryPaginator(@Nonnull Query query, int pageSize) { + this.hasOverallLimit = query.builder.getRowsLimit() == 0 ? false : true; + this.remainingRows = query.builder.getRowsLimit(); this.query = query; - this.chunkSize = chunkSize; - this.prevSplitPoint = ByteString.EMPTY; - this.firstRun = true; + this.pageSize = pageSize; + this.prevSplitPoint = null; } /** Return the next query. Needs to be called after advance(). */ - Query getNextQuery() { + public Query getNextQuery() { return query; } @@ -353,30 +351,26 @@ Query getNextQuery() { * Construct the next query. Return true if there are more queries to return. False if we've * read everything. */ - boolean advance(@Nonnull ByteString lastSeenRowKey) { - Preconditions.checkArgument( - lastSeenRowKey != null, "lastSeenRowKey cannot be null, use ByteString.EMPTY instead."); + public boolean advance(@Nonnull ByteString lastSeenRowKey) { + Preconditions.checkNotNull( + lastSeenRowKey, "lastSeenRowKey cannot be null, use ByteString.EMPTY instead."); // Full table scans don't have ranges or limits. Keep track of the previous split - // point. If it's the same as the current input return false. The only exception - // is the first run for this paginator. So keep track of the first run state too. - if (!firstRun && prevSplitPoint.equals(lastSeenRowKey)) { + // point. If it's the same as the current input return false. + if (lastSeenRowKey.equals(prevSplitPoint)) { return false; } - if (firstRun) { - firstRun = false; - } this.prevSplitPoint = lastSeenRowKey; // Set the query limit. If the original limit is set, return false if the new // limit is <= 0 to avoid returning more rows than intended. - if (originalLimit != 0 && newLimit <= 0) { + if (hasOverallLimit && remainingRows <= 0) { return false; } - if (originalLimit != 0) { - query.limit(Math.min(this.chunkSize, newLimit)); - newLimit -= chunkSize; + if (hasOverallLimit) { + query.limit(Math.min(this.pageSize, remainingRows)); + remainingRows -= pageSize; } else { - query.limit(chunkSize); + query.limit(pageSize); } // Split the row ranges / row keys. Return false if there's nothing diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index 79fe4a32fe..f144863817 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -332,7 +332,7 @@ public void testClone() { public void testQueryPaginatorRangeLimitReached() { int chunkSize = 10, limit = 15; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); - Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator paginator = query.createPaginator(chunkSize); assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -371,7 +371,7 @@ public void testQueryPaginatorRangeLimitReached() { public void testQueryPaginatorRangeLimitMultiplyOfChunkSize() { int chunkSize = 10, limit = 20; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); - Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator paginator = query.createPaginator(chunkSize); assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -410,7 +410,7 @@ public void testQueryPaginatorRangeLimitMultiplyOfChunkSize() { public void testQueryPaginatorRagneNoLimit() { int chunkSize = 10; Query query = Query.create(TABLE_ID).range("a", "z"); - Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator paginator = query.createPaginator(chunkSize); assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -448,7 +448,7 @@ public void testQueryPaginatorRowsNoLimit() { int chunkSize = 10; Query query = Query.create(TABLE_ID).rowKey("a").rowKey("b").rowKey("c"); - Query.QueryPaginator paginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator paginator = query.createPaginator(chunkSize); assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -478,7 +478,7 @@ public void testQueryPaginatorRowsNoLimit() { public void testQueryPaginatorFullTableScan() { int chunkSize = 10; Query query = Query.create(TABLE_ID); - Query.QueryPaginator queryPaginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue(); ReadRowsRequest.Builder expectedProto = @@ -503,7 +503,7 @@ public void testQueryPaginatorFullTableScan() { public void testQueryPaginatorEmptyTable() { int chunkSize = 10; Query query = Query.create(TABLE_ID); - Query.QueryPaginator queryPaginator = query.createQueryPaginator(chunkSize); + Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue(); ReadRowsRequest.Builder expectedProto = From 279f58aa29bc706e151d2ca62df8a98eab333c5e Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 29 Nov 2022 15:52:58 -0500 Subject: [PATCH 6/9] update --- .../google/cloud/bigtable/data/v2/models/Query.java | 13 +++++++------ .../cloud/bigtable/data/v2/models/QueryTest.java | 6 ------ 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index 6d9880c3e4..1149760dc3 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -258,13 +258,13 @@ public List shard(SortedSet splitPoints) { * Query query = Query.create(...).range("a", "z"); * Query.QueryPaginator paginator = query.createQueryPaginator(100); * ByteString lastSeenRowKey = ByteString.EMPTY; - * while (paginator.advance(lastSeenRowKey)) { + * do { * List rows = client.readRowsCallable().all().call(paginator.getNextQuery()); * for (Row row : rows) { * // do some processing * lastSeenRow = row; * } - * } + * } while (paginator.advance(lastSeenRowKey)); * } */ @BetaApi("This surface is stable yet it might be removed in the future.") @@ -337,12 +337,12 @@ public static class QueryPaginator { QueryPaginator(@Nonnull Query query, int pageSize) { this.hasOverallLimit = query.builder.getRowsLimit() == 0 ? false : true; this.remainingRows = query.builder.getRowsLimit(); - this.query = query; + this.query = query.limit(pageSize); this.pageSize = pageSize; this.prevSplitPoint = null; } - /** Return the next query. Needs to be called after advance(). */ + /** Return the next query. */ public Query getNextQuery() { return query; } @@ -354,8 +354,9 @@ public Query getNextQuery() { public boolean advance(@Nonnull ByteString lastSeenRowKey) { Preconditions.checkNotNull( lastSeenRowKey, "lastSeenRowKey cannot be null, use ByteString.EMPTY instead."); - // Full table scans don't have ranges or limits. Keep track of the previous split - // point. If it's the same as the current input return false. + // Full table scans don't have ranges or limits. Running the query again will return an empty + // list when we reach the end of the table. lastSeenRowKey won't be updated in this case, and + // we can break out of the loop. if (lastSeenRowKey.equals(prevSplitPoint)) { return false; } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index f144863817..35df648b66 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -333,7 +333,6 @@ public void testQueryPaginatorRangeLimitReached() { int chunkSize = 10, limit = 15; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); Query.QueryPaginator paginator = query.createPaginator(chunkSize); - assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -372,7 +371,6 @@ public void testQueryPaginatorRangeLimitMultiplyOfChunkSize() { int chunkSize = 10, limit = 20; Query query = Query.create(TABLE_ID).range("a", "z").limit(limit); Query.QueryPaginator paginator = query.createPaginator(chunkSize); - assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -411,7 +409,6 @@ public void testQueryPaginatorRagneNoLimit() { int chunkSize = 10; Query query = Query.create(TABLE_ID).range("a", "z"); Query.QueryPaginator paginator = query.createPaginator(chunkSize); - assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -449,7 +446,6 @@ public void testQueryPaginatorRowsNoLimit() { Query query = Query.create(TABLE_ID).rowKey("a").rowKey("b").rowKey("c"); Query.QueryPaginator paginator = query.createPaginator(chunkSize); - assertThat(paginator.advance(ByteString.EMPTY)).isTrue(); Query nextQuery = paginator.getNextQuery(); @@ -479,7 +475,6 @@ public void testQueryPaginatorFullTableScan() { int chunkSize = 10; Query query = Query.create(TABLE_ID); Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); - assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue(); ReadRowsRequest.Builder expectedProto = expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize); @@ -504,7 +499,6 @@ public void testQueryPaginatorEmptyTable() { int chunkSize = 10; Query query = Query.create(TABLE_ID); Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); - assertThat(queryPaginator.advance(ByteString.EMPTY)).isTrue(); ReadRowsRequest.Builder expectedProto = expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize); From ff7ed8b648b13a6c2deca554e583ccf73672e304 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 29 Nov 2022 16:03:13 -0500 Subject: [PATCH 7/9] fix test --- .../com/google/cloud/bigtable/data/v2/models/Query.java | 5 ++++- .../com/google/cloud/bigtable/data/v2/models/QueryTest.java | 6 ++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index 1149760dc3..abb46d46d9 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -338,8 +338,11 @@ public static class QueryPaginator { this.hasOverallLimit = query.builder.getRowsLimit() == 0 ? false : true; this.remainingRows = query.builder.getRowsLimit(); this.query = query.limit(pageSize); + if (hasOverallLimit) { + remainingRows -= pageSize; + } this.pageSize = pageSize; - this.prevSplitPoint = null; + this.prevSplitPoint = ByteString.EMPTY; } /** Return the next query. */ diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java index 35df648b66..655aeda688 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/QueryTest.java @@ -476,8 +476,7 @@ public void testQueryPaginatorFullTableScan() { Query query = Query.create(TABLE_ID); Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); - ReadRowsRequest.Builder expectedProto = - expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize); + ReadRowsRequest.Builder expectedProto = expectedProtoBuilder().setRowsLimit(chunkSize); assertThat(queryPaginator.getNextQuery().toProto(requestContext)) .isEqualTo(expectedProto.build()); @@ -500,8 +499,7 @@ public void testQueryPaginatorEmptyTable() { Query query = Query.create(TABLE_ID); Query.QueryPaginator queryPaginator = query.createPaginator(chunkSize); - ReadRowsRequest.Builder expectedProto = - expectedProtoBuilder().setRows(RowSet.getDefaultInstance()).setRowsLimit(chunkSize); + ReadRowsRequest.Builder expectedProto = expectedProtoBuilder().setRowsLimit(chunkSize); assertThat(queryPaginator.getNextQuery().toProto(requestContext)) .isEqualTo(expectedProto.build()); From d11565541cac606626995f8d41aff120a8f14c6f Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 29 Nov 2022 16:13:59 -0500 Subject: [PATCH 8/9] fix nit --- .../java/com/google/cloud/bigtable/data/v2/models/Query.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java index abb46d46d9..271ffe3adf 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Query.java @@ -335,7 +335,7 @@ public static class QueryPaginator { private ByteString prevSplitPoint; QueryPaginator(@Nonnull Query query, int pageSize) { - this.hasOverallLimit = query.builder.getRowsLimit() == 0 ? false : true; + this.hasOverallLimit = query.builder.getRowsLimit() > 0; this.remainingRows = query.builder.getRowsLimit(); this.query = query.limit(pageSize); if (hasOverallLimit) { From 0860a38166cae22f391179b36e14cb37468df2f4 Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Tue, 29 Nov 2022 21:20:01 +0000 Subject: [PATCH 9/9] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20po?= =?UTF-8?q?st-processor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 71f9e3fcd5..5ce31d082e 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ If you are using Maven without BOM, add this to your dependencies: If you are using Gradle 5.x or later, add this to your dependencies: ```Groovy -implementation platform('com.google.cloud:libraries-bom:26.1.4') +implementation platform('com.google.cloud:libraries-bom:26.1.5') implementation 'com.google.cloud:google-cloud-bigtable' ```