fix(bigquery): preserve page token and paginate correctly in Arrow query when maxResults is set (#14469)

When `maxResults` is set on `QueryJobConfiguration` for row-based Arrow
queries, the initial REST response uses it as the page size for page 1.
Previously, the client treated `initialRowOffset >= maxResults` as an
indication to terminate pagination early and passed `maxResults` as the
total row ceiling to `ArrowQueryPageFetcher`, preventing subsequent
pages from being fetched.

This updates `queryRpcArrow` to preserve the page token, pass
`maxResults` as the page size option to `ArrowQueryPageFetcher`, and
verifies pagination behavior in both unit tests and
`testQueryRowBasedWithArrowFormatMultiPage`.

Author: jinseopkim0

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -74,6 +74,7 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -2655,11 +2656,6 @@ && getOptions().getOpenTelemetryTracer() != null) {
     // Calculate row offset and determine if subsequent pages exist.
     boolean hasMorePages = results.getPageToken() != null;
     long initialRowOffset = (long) firstPageRows.size();
-    if (hasMorePages) {
-      if (content.getMaxResults() != null && initialRowOffset >= content.getMaxResults()) {
-        hasMorePages = false;
-      }
-    }
 
     // Multi-page results: configure ArrowQueryPageFetcher for subsequent tabledata.list calls.
     if (hasMorePages) {
@@ -2669,6 +2665,11 @@ && getOptions().getOpenTelemetryTracer() != null) {
       }
       JobId jobId = JobId.fromPb(results.getJobReference());
       String cursor = results.getPageToken();
+      Map<BigQueryRpc.Option, Object> fetcherOptions = new HashMap<>(optionMap(options));
+      if (content.getMaxResults() != null
+          && !fetcherOptions.containsKey(BigQueryRpc.Option.MAX_RESULTS)) {
+        fetcherOptions.put(BigQueryRpc.Option.MAX_RESULTS, content.getMaxResults());
+      }
       NextPageFetcher<FieldValueList> pageFetcher =
           new ArrowQueryPageFetcher(
               jobId,
@@ -2677,8 +2678,8 @@ && getOptions().getOpenTelemetryTracer() != null) {
               arrowSchemaPojo,
               getOptions(),
               initialRowOffset,
-              content.getMaxResults(),
-              optionMap(options));
+              null,
+              fetcherOptions);
 
       return newTableResultBuilder(results)
           .setSchema(schema)
@@ -3190,6 +3191,10 @@ private TableResult readArrowTableResultFromJob(
     }
 
     // Initialize the page fetcher targeting the ReadSession stream to load the first page of rows.
+    Map<BigQueryRpc.Option, Object> fetcherOptions = new HashMap<>(optionMap(options));
+    if (maxResults != null && !fetcherOptions.containsKey(BigQueryRpc.Option.MAX_RESULTS)) {
+      fetcherOptions.put(BigQueryRpc.Option.MAX_RESULTS, maxResults);
+    }
     ArrowQueryPageFetcher pageFetcher =
         new ArrowQueryPageFetcher(
             completedJob.getJobId(),
@@ -3199,8 +3204,8 @@ private TableResult readArrowTableResultFromJob(
             arrowSchemaPojo,
             getOptions(),
             0L,
-            maxResults,
-            optionMap(options));
+            null,
+            fetcherOptions);
 
     Page<FieldValueList> firstPage = pageFetcher.getNextPage();
     List<FieldValueList> firstPageRows =

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -3354,7 +3354,8 @@ void testQueryWithArrowFormatMultiplePagesWithMaxResults()
             .build();
     ReadRowsResponse streamResponse =
         ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build();
-    when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator());
+    when(mockServerStream.iterator())
+        .thenAnswer(invocation -> ImmutableList.of(streamResponse).iterator());
 
     BigQueryReadClient mockReadClient =
         mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
@@ -3375,25 +3376,28 @@ void testQueryWithArrowFormatMultiplePagesWithMaxResults()
     Page<FieldValueList> page2 = result.getNextPage();
     assertNotNull(page2);
     List<FieldValueList> page2Rows = ImmutableList.copyOf(page2.getValues());
-    // Since maxResults is 2 and initialRowOffset is 1, page2 should only contain 1 row even though
-    // stream returned 2 rows
-    assertEquals(1, page2Rows.size());
+    // Since maxResults configures the page size (2 rows), page2 contains the 2 rows from the stream
+    assertEquals(2, page2Rows.size());
     assertEquals("2", page2Rows.get(0).get(0).getStringValue());
-    // Since totalRowsReturned == maxResults, hasNextPage must be false
+    assertEquals("3", page2Rows.get(1).get(0).getStringValue());
+    // End of stream reached (total 3 rows read across pages 1 and 2), hasNextPage must be false
     assertFalse(page2.hasNextPage());
     assertNull(page2.getNextPage());
 
-    // When maxResults is 1, initialRowOffset (1) already reaches maxResults, so hasNextPage is
-    // false immediately
+    // When maxResults is 1, page token is still preserved for subsequent pages
     QueryJobConfiguration configMax1 =
         QueryJobConfiguration.newBuilder("SELECT id FROM test")
             .setQueryResultsFormat(QueryResultsFormat.ARROW)
             .setMaxResults(1L)
             .build();
     TableResult resultMax1 = bigquery.query(configMax1);
     assertNotNull(resultMax1);
-    assertFalse(resultMax1.hasNextPage());
-    assertNull(resultMax1.getNextPage());
+    assertTrue(resultMax1.hasNextPage());
+    Page<FieldValueList> page2Max1 = resultMax1.getNextPage();
+    assertNotNull(page2Max1);
+    List<FieldValueList> page2Max1Rows = ImmutableList.copyOf(page2Max1.getValues());
+    assertEquals(1, page2Max1Rows.size());
+    assertEquals("2", page2Max1Rows.get(0).get(0).getStringValue());
   }
 
   @Test

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -7570,20 +7570,32 @@ void testQueryRowBasedWithArrowFormat() throws InterruptedException {
 
   @Test
   void testQueryRowBasedWithArrowFormatMultiPage() throws InterruptedException {
+    // Under fast-query execution, the initial REST response defaults to a 10 MB payload limit.
+    // Setting maxResults limits the initial page to 5,000 rows, forcing the remaining 10,000 rows
+    // to stream across multiple pages via tabledata.list with Arrow format.
     String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x";
     QueryJobConfiguration config =
         QueryJobConfiguration.newBuilder(query)
             .setQueryResultsFormat(QueryResultsFormat.ARROW)
             .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
+            .setMaxResults(5000L)
             .build();
     TableResult result = bigquery.query(config);
     assertNotNull(result);
     assertEquals(15000, result.getTotalRows());
+
+    int pageCount = 0;
     long count = 0;
-    for (FieldValueList row : result.iterateAll()) {
-      count++;
-      assertEquals(count, row.get("x").getLongValue());
+    TableResult currentPage = result;
+    while (currentPage != null) {
+      pageCount++;
+      for (FieldValueList row : currentPage.getValues()) {
+        count++;
+        assertEquals(count, row.get("x").getLongValue());
+      }
+      currentPage = currentPage.hasNextPage() ? currentPage.getNextPage() : null;
     }
+    assertTrue(pageCount > 1);
     assertEquals(15000, count);
   }
 
@@ -7635,6 +7647,62 @@ void testQueryRowBasedWithArrowFormatFallback() throws InterruptedException {
     assertEquals(1786363200000000L, row.get("ts").getTimestampValue());
   }
 
+  @Test
+  void testQueryResultsFormatArrowFallbackMultiPage() throws InterruptedException {
+    String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x";
+    QueryJobConfiguration config =
+        QueryJobConfiguration.newBuilder(query)
+            .setQueryResultsFormat(QueryResultsFormat.ARROW)
+            .setMaxResults(5000L)
+            .build();
+    JobId customJobId =
+        JobId.of("arrow_it_fallback_mp_" + UUID.randomUUID().toString().replace("-", "_"));
+    try (ArrowQueryResult result = bigquery.queryArrow(config, customJobId)) {
+      assertNotNull(result);
+      assertNotNull(result.getJobId());
+      assertEquals(customJobId.getJob(), result.getJobId().getJob());
+      int batchCount = 0;
+      long totalRows = 0;
+      for (VectorSchemaRoot root : result) {
+        batchCount++;
+        totalRows += root.getRowCount();
+      }
+      assertTrue(batchCount > 1);
+      assertEquals(15000, totalRows);
+    }
+  }
+
+  @Test
+  void testQueryRowBasedWithArrowFormatFallbackMultiPage() throws InterruptedException {
+    String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x";
+    QueryJobConfiguration config =
+        QueryJobConfiguration.newBuilder(query)
+            .setQueryResultsFormat(QueryResultsFormat.ARROW)
+            .setMaxResults(5000L)
+            .build();
+    JobId customJobId =
+        JobId.of("row_it_fallback_mp_" + UUID.randomUUID().toString().replace("-", "_"));
+    TableResult result = bigquery.query(config, customJobId);
+    assertNotNull(result);
+    assertNotNull(result.getJobId());
+    assertEquals(customJobId.getJob(), result.getJobId().getJob());
+    assertEquals(15000, result.getTotalRows());
+
+    int pageCount = 0;
+    long count = 0;
+    TableResult currentPage = result;
+    while (currentPage != null) {
+      pageCount++;
+      for (FieldValueList row : currentPage.getValues()) {
+        count++;
+        assertEquals(count, row.get("x").getLongValue());
+      }
+      currentPage = currentPage.hasNextPage() ? currentPage.getNextPage() : null;
+    }
+    assertTrue(pageCount > 1);
+    assertEquals(15000, count);
+  }
+
   @Test
   void testUniverseDomainWithInvalidUniverseDomain() {
     RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();