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

Author: jinseopkim0

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

@@ -2620,11 +2620,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) {
@@ -2634,6 +2629,11 @@ && getOptions().getOpenTelemetryTracer() != null) {
       }
       JobId jobId = JobId.fromPb(results.getJobReference());
       String cursor = results.getPageToken();
+      Map<BigQueryRpc.Option, Object> fetcherOptions = new java.util.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,
@@ -2642,8 +2642,8 @@ && getOptions().getOpenTelemetryTracer() != null) {
               arrowSchemaPojo,
               getOptions(),
               initialRowOffset,
-              content.getMaxResults(),
-              optionMap(options));
+              null,
+              fetcherOptions);
 
       return newTableResultBuilder(results)
           .setSchema(schema)

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