Skip to content

Commit 615409f

Browse files
authored
feat(bigquery): add ArrowQueryPageFetcher for Arrow query result pagination (#14404)
This PR introduces `ArrowQueryPageFetcher` to handle pagination for queries executing with `QueryResultsFormat.ARROW`. ### Summary of Changes - Implements `ArrowQueryPageFetcher` implementing `NextPageFetcher<FieldValueList>`. - Connects to the default storage read stream to fetch subsequent row pages. - Leverages `ArrowDeserializer.loadArrowRows` to deserialize Arrow record batches into `FieldValueList` collections with offset and `maxResults` bounding. - Adds `ArrowQueryPageFetcherTest` covering single-page, multi-page, max-results, and serialization behaviors.
1 parent 4505e95 commit 615409f

2 files changed

Lines changed: 568 additions & 0 deletions

File tree

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

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.api.gax.paging.Page;
2929
import com.google.api.gax.rpc.HeaderProvider;
3030
import com.google.api.gax.rpc.NoHeaderProvider;
31+
import com.google.api.gax.rpc.ServerStream;
3132
import com.google.api.services.bigquery.model.ErrorProto;
3233
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
3334
import com.google.api.services.bigquery.model.ProjectList;
@@ -54,6 +55,8 @@
5455
import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings;
5556
import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
5657
import com.google.cloud.bigquery.storage.v1.DataFormat;
58+
import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
59+
import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
5760
import com.google.cloud.bigquery.storage.v1.ReadSession;
5861
import com.google.common.annotations.VisibleForTesting;
5962
import com.google.common.base.Function;
@@ -68,10 +71,13 @@
6871
import io.opentelemetry.api.trace.Span;
6972
import io.opentelemetry.context.Scope;
7073
import java.io.IOException;
74+
import java.util.ArrayDeque;
7175
import java.util.ArrayList;
7276
import java.util.Collections;
77+
import java.util.Iterator;
7378
import java.util.List;
7479
import java.util.Map;
80+
import java.util.Queue;
7581
import java.util.concurrent.Callable;
7682
import java.util.concurrent.ConcurrentHashMap;
7783
import java.util.regex.Matcher;
@@ -278,6 +284,192 @@ public Page<FieldValueList> getNextPage() {
278284
}
279285
}
280286

287+
/**
288+
* NextPageFetcher implementation for queries returning results in Arrow format. Reads subsequent
289+
* pages from the job's default gRPC storage read stream. Rows are buffered and converted on
290+
* demand from Arrow record batches into FieldValueList instances in batches up to {@code
291+
* pageSize}.
292+
*
293+
* <p>Note: Neither {@link Page} nor {@link TableResult} implements {@link AutoCloseable}. The
294+
* underlying gRPC stream is automatically canceled and resources released when iteration reaches
295+
* the end (or maximum results requested) or when an error occurs. Callers that do not iterate to
296+
* completion rely on server-side stream timeouts and garbage collection to release stream
297+
* resources.
298+
*/
299+
static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
300+
private static final long serialVersionUID = 1L;
301+
private static final long DEFAULT_PAGE_SIZE = 10000L;
302+
303+
private final JobId jobId;
304+
private final Schema schema;
305+
private final byte[] arrowSchemaBytes;
306+
private final BigQueryOptions serviceOptions;
307+
private final long maxResults;
308+
private final Map<BigQueryRpc.Option, ?> optionsMap;
309+
310+
private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo;
311+
private transient BigQueryReadClient bqReadClient;
312+
private transient ServerStream<ReadRowsResponse> stream;
313+
private transient Iterator<ReadRowsResponse> streamIterator;
314+
private transient Queue<FieldValueList> buffer = new ArrayDeque<>();
315+
private long totalRowsReturned = 0L;
316+
private boolean streamClosed = false;
317+
318+
ArrowQueryPageFetcher(
319+
JobId jobId,
320+
Schema schema,
321+
byte[] arrowSchemaBytes,
322+
org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo,
323+
BigQueryOptions serviceOptions,
324+
long initialRowOffset,
325+
Long maxResults,
326+
Map<BigQueryRpc.Option, ?> optionsMap) {
327+
this.jobId = jobId;
328+
this.schema = schema;
329+
this.arrowSchemaBytes = arrowSchemaBytes;
330+
this.arrowSchemaPojo = arrowSchemaPojo;
331+
this.serviceOptions = serviceOptions;
332+
this.totalRowsReturned = initialRowOffset;
333+
this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE;
334+
this.optionsMap = optionsMap;
335+
}
336+
337+
@Override
338+
public Page<FieldValueList> getNextPage() {
339+
if (buffer == null) {
340+
buffer = new ArrayDeque<>();
341+
}
342+
if (streamClosed || totalRowsReturned >= maxResults) {
343+
closeClient();
344+
return null;
345+
}
346+
347+
Number optionPageSize =
348+
optionsMap != null ? (Number) optionsMap.get(BigQueryRpc.Option.MAX_RESULTS) : null;
349+
long pageSize =
350+
optionPageSize != null && optionPageSize.longValue() > 0
351+
? optionPageSize.longValue()
352+
: DEFAULT_PAGE_SIZE;
353+
List<FieldValueList> rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L));
354+
355+
try {
356+
// Resolve job location in order: JobId location -> BigQueryOptions location -> "global"
357+
// default.
358+
// The Storage Read API stream resource name requires a location component (e.g.
359+
// projects/{project}/locations/{location}/jobs/{job}/streams/_default). If no specific
360+
// location
361+
// was provided on the job or service options, defaulting to "global" allows queries created
362+
// without an explicit location to still stream results without failing.
363+
String location = jobId.getLocation();
364+
if (location == null) {
365+
location = serviceOptions.getLocation();
366+
}
367+
if (location == null) {
368+
location = "global";
369+
}
370+
371+
if (streamIterator == null) {
372+
if (bqReadClient == null) {
373+
BigQuery service = serviceOptions.getService();
374+
if (!(service instanceof BigQueryImpl)) {
375+
// Arrow pagination relies on BigQueryImpl to manage and cache the underlying
376+
// BigQueryReadClient across page fetches. Custom BigQuery implementations that do
377+
// not extend BigQueryImpl cannot provide the managed gRPC storage read client.
378+
throw new IllegalStateException(
379+
"Arrow query result pagination requires an instance of BigQueryImpl to manage BigQueryReadClient lifecycle");
380+
}
381+
bqReadClient = ((BigQueryImpl) service).getBigQueryReadClient(location);
382+
}
383+
384+
// Construct the default stream path for reading job query results via Storage Read API.
385+
String streamName =
386+
String.format(
387+
"projects/%s/locations/%s/jobs/%s/streams/_default",
388+
jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
389+
location,
390+
jobId.getJob());
391+
392+
ReadRowsRequest readRowsRequest =
393+
ReadRowsRequest.newBuilder()
394+
.setReadStream(streamName)
395+
.setOffset(totalRowsReturned)
396+
.build();
397+
398+
// Open the server-streaming RPC and obtain the response iterator.
399+
stream = bqReadClient.readRowsCallable().call(readRowsRequest);
400+
streamIterator = stream.iterator();
401+
}
402+
403+
// Lazily deserialize the Arrow schema from serialized bytes on the first page fetch (e.g.,
404+
// after fetcher deserialization) and cache the POJO schema for subsequent batches.
405+
if (arrowSchemaPojo == null && arrowSchemaBytes != null) {
406+
arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
407+
}
408+
409+
boolean hasMore =
410+
ArrowDeserializer.loadArrowRows(
411+
streamIterator,
412+
arrowSchemaPojo,
413+
schema,
414+
rowBatch,
415+
buffer,
416+
pageSize,
417+
totalRowsReturned,
418+
maxResults);
419+
420+
// If no rows were read (e.g., EOF reached on an empty stream or after all results
421+
// were exhausted), terminate pagination. Note: when hasMore is false, rowBatch may
422+
// still contain the final page of data to return, so we only return null if rowBatch
423+
// is empty.
424+
if (rowBatch.isEmpty()) {
425+
streamClosed = true;
426+
closeClient();
427+
return null;
428+
}
429+
430+
totalRowsReturned += rowBatch.size();
431+
432+
// Generate nextPageToken only if there are more rows in the stream and maxResults
433+
// has not been reached. Otherwise, close the stream early to free resources.
434+
String nextPageToken = null;
435+
if (hasMore && totalRowsReturned < maxResults) {
436+
nextPageToken = String.valueOf(totalRowsReturned);
437+
} else {
438+
streamClosed = true;
439+
closeClient();
440+
}
441+
442+
return new PageImpl<>(this, nextPageToken, rowBatch);
443+
444+
} catch (BigQueryException e) {
445+
streamClosed = true;
446+
closeClient();
447+
throw e;
448+
} catch (Exception e) {
449+
streamClosed = true;
450+
closeClient();
451+
throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e);
452+
}
453+
}
454+
455+
/**
456+
* Cancels the active Storage Read API server stream (if any) and clears transient references so
457+
* gRPC stream resources can be reclaimed immediately.
458+
*/
459+
private void closeClient() {
460+
if (stream != null) {
461+
try {
462+
stream.cancel();
463+
} catch (Exception e) {
464+
// Ignore cancellation exceptions during teardown
465+
}
466+
}
467+
bqReadClient = null;
468+
streamIterator = null;
469+
stream = null;
470+
}
471+
}
472+
281473
private static final int MAX_CACHED_READ_CLIENTS = 100;
282474

283475
private transient ConcurrentHashMap<String, BigQueryReadClient> bqReadClients;

0 commit comments

Comments
 (0)