udpate test

Change-Id: I8bd49e79fb8d9c368a5126e7fa0bb4f8d73cff3c

Author: mutianf

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

@@ -17,62 +17,68 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.bigtable.data.v2.models.Filters;
+import javax.annotation.Nullable;
 
 /**
- * A paginator for fetching extremely large rows from Bigtable chunk by chunk to avoid the 256MB gRPC size limit.
- * It yields Filters that chunk the row by cell limits and handles limit halving if FAILED_PRECONDITION occurs.
+ * A paginator for fetching large rows from Bigtable chunk by chunk to avoid the 256MB size limit.
+ * It yields Filters that chunk the row by cell limits and handles limit halving if
+ * FAILED_PRECONDITION occurs.
  */
 @InternalApi("For internal usage only")
 public class LargeRowPaginator {
-    private int currentLimit;
-    private int currentOffset;
-    private boolean hasMore;
+  private int currentLimit;
+  private int currentOffset;
+  private boolean hasMore;
+  @Nullable private final Filters.Filter baseFilter;
 
-    public LargeRowPaginator(int initialLimit) {
-        this.currentLimit = initialLimit;
-        this.currentOffset = 0;
-        this.hasMore = true;
-    }
+  public LargeRowPaginator(int initialLimit, @Nullable Filters.Filter filter) {
+    this.currentLimit = initialLimit;
+    this.currentOffset = 0;
+    this.hasMore = true;
+    this.baseFilter = filter;
+  }
 
-    /**
-     * Yields the filter required to fetch the next chunk of cells for the large row.
-     */
-    public Filters.Filter getNextFilter() {
-        Filters.ChainFilter chain = Filters.FILTERS.chain();
-        if (currentOffset > 0) {
-            chain.filter(Filters.FILTERS.offset().cellsPerRow(currentOffset));
-        }
-        chain.filter(Filters.FILTERS.limit().cellsPerRow(currentLimit));
-        return chain;
+  /** Yields the filter required to fetch the next chunk of cells for the large row. */
+  public Filters.Filter getNextFilter() {
+    Filters.ChainFilter chain = Filters.FILTERS.chain();
+    if (baseFilter != null) {
+      chain.filter(baseFilter);
     }
-
-    /**
-     * Advances the internal offset. Call this after a successful Bigtable API call.
-     * @param cellsReadInLastChunk The number of cells returned in the last chunk.
-     * @return true if there are potentially more cells to fetch.
-     */
-    public boolean advance(int cellsReadInLastChunk) {
-        this.currentOffset += cellsReadInLastChunk;
-        
-        // If we read fewer cells than requested, we've hit the end of the row.
-        if (cellsReadInLastChunk < currentLimit) {
-            this.hasMore = false;
-        }
-        return this.hasMore;
+    if (currentOffset > 0) {
+      chain.filter(Filters.FILTERS.offset().cellsPerRow(currentOffset));
     }
+    chain.filter(Filters.FILTERS.limit().cellsPerRow(currentLimit));
+    return chain;
+  }
 
-    /**
-     * Call this if the Bigtable API call fails with a FAILED_PRECONDITION due to size limits.
-     * It reduces the batch size to fetch a smaller chunk on the next attempt.
-     */
-    public void halveLimit() {
-        this.currentLimit /= 2;
-        if (this.currentLimit == 0) {
-            throw new RuntimeException("Cannot divide limit further. A single cell might be too large.");
-        }
+  /**
+   * Advances the internal offset. Call this after a successful Bigtable API call.
+   *
+   * @param cellsReadInLastChunk The number of cells returned in the last chunk.
+   * @return true if there are potentially more cells to fetch.
+   */
+  public boolean advance(int cellsReadInLastChunk) {
+    this.currentOffset += cellsReadInLastChunk;
+
+    // If we read fewer cells than requested, we've hit the end of the row.
+    if (cellsReadInLastChunk < currentLimit) {
+      this.hasMore = false;
     }
+    return this.hasMore;
+  }
 
-    public boolean hasNext() {
-        return this.hasMore;
+  /**
+   * Call this if the Bigtable API call fails with a FAILED_PRECONDITION due to size limits. It
+   * reduces the batch size to fetch a smaller chunk on the next attempt.
+   */
+  public void halveLimit() {
+    this.currentLimit /= 2;
+    if (this.currentLimit == 0) {
+      throw new RuntimeException("Cannot divide limit further. A single cell might be too large.");
     }
+  }
+
+  public boolean hasNext() {
+    return this.hasMore;
+  }
 }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/LargeRowPaginatorTest.java

@@ -18,6 +18,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 
+import com.google.cloud.bigtable.data.v2.models.Filters;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -27,15 +28,15 @@ public class LargeRowPaginatorTest {
 
   @Test
   public void testPaginatorAdvancesProperly() {
-    LargeRowPaginator paginator = new LargeRowPaginator(10);
-    
+    LargeRowPaginator paginator = new LargeRowPaginator(10, null);
+
     assertThat(paginator.hasNext()).isTrue();
-    
+
     // Simulate reading exactly the limit (10 cells). Paginator assumes more data exists.
     boolean hasMore = paginator.advance(10);
     assertThat(hasMore).isTrue();
     assertThat(paginator.hasNext()).isTrue();
-    
+
     // Simulate reading 5 cells (less than limit, indicating end of row)
     hasMore = paginator.advance(5);
     assertThat(hasMore).isFalse();
@@ -44,26 +45,42 @@ public void testPaginatorAdvancesProperly() {
 
   @Test
   public void testPaginatorHalvesLimit() {
-    LargeRowPaginator paginator = new LargeRowPaginator(10);
-    
+    LargeRowPaginator paginator = new LargeRowPaginator(10, null);
+
     paginator.halveLimit(); // Internal limit becomes 5
-    
+
     // Simulate reading exactly the new limit (5 cells)
     boolean hasMore = paginator.advance(5);
-    assertThat(hasMore).isTrue(); 
-    
+    assertThat(hasMore).isTrue();
+
     paginator.halveLimit(); // Internal limit becomes 2
-    
+
     // Simulate reading 1 cell (less than the new limit of 2)
-    hasMore = paginator.advance(1); 
+    hasMore = paginator.advance(1);
     assertThat(hasMore).isFalse();
   }
 
   @Test
   public void testPaginatorThrowsOnZeroLimit() {
-    LargeRowPaginator paginator = new LargeRowPaginator(1);
-    
+    LargeRowPaginator paginator = new LargeRowPaginator(1, null);
+
     RuntimeException exception = assertThrows(RuntimeException.class, () -> paginator.halveLimit());
     assertThat(exception).hasMessageThat().contains("Cannot divide limit further");
   }
+
+  @Test
+  public void testPaginatorWithBaseFilter() {
+    LargeRowPaginator paginator =
+        new LargeRowPaginator(10, Filters.FILTERS.family().exactMatch("cf"));
+
+    Filters.Filter nextFilter = paginator.getNextFilter();
+
+    Filters.Filter expectedFilter =
+        Filters.FILTERS
+            .chain()
+            .filter(Filters.FILTERS.family().exactMatch("cf"))
+            .filter(Filters.FILTERS.limit().cellsPerRow(10));
+
+    assertThat(nextFilter.toProto()).isEqualTo(expectedFilter.toProto());
+  }
 }