Skip to content

feat(bigquery): add Storage Read API slow-path fallback for row-based query() #14389

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jinseopkim0 wants to merge 1 commit into from
+328 −25
Draft

feat(bigquery): add Storage Read API slow-path fallback for row-based query() #14389

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
223 changes: 208 additions & 15 deletions ...-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,7 @@ static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
private static final long DEFAULT_PAGE_SIZE = 10000L;

private final JobId jobId;
private final String customStreamName;
private final Schema schema;
private final byte[] arrowSchemaBytes;
private final BigQueryOptions serviceOptions;
Expand All @@ -327,7 +328,30 @@ static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
long initialRowOffset,
Long maxResults,
Map<BigQueryRpc.Option, ?> optionsMap) {
this(
jobId,
null,
schema,
arrowSchemaBytes,
arrowSchemaPojo,
serviceOptions,
initialRowOffset,
maxResults,
optionsMap);
}

ArrowQueryPageFetcher(
JobId jobId,
String customStreamName,
Schema schema,
byte[] arrowSchemaBytes,
org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo,
BigQueryOptions serviceOptions,
long initialRowOffset,
Long maxResults,
Map<BigQueryRpc.Option, ?> optionsMap) {
this.jobId = jobId;
this.customStreamName = customStreamName;
this.schema = schema;
this.arrowSchemaBytes = arrowSchemaBytes;
this.arrowSchemaPojo = arrowSchemaPojo;
Expand Down Expand Up @@ -356,13 +380,6 @@ public Page<FieldValueList> getNextPage() {
List<FieldValueList> rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L));

try {
String location =
jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
if (location == null) {
throw new BigQueryException(
0, "Location must be specified to read Arrow rows from storage stream");
}

if (bqReadClient == null) {
BigQuery service = serviceOptions.getService();
if (service instanceof BigQueryImpl) {
Expand All @@ -375,12 +392,23 @@ public Page<FieldValueList> getNextPage() {
}

if (streamIterator == null) {
String streamName =
String.format(
"projects/%s/locations/%s/jobs/%s/streams/_default",
jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
location,
jobId.getJob());
String streamName;
if (customStreamName != null) {
streamName = customStreamName;
} else {
String location =
jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
if (location == null) {
throw new BigQueryException(
0, "Location must be specified to read Arrow rows from storage stream");
}
streamName =
String.format(
"projects/%s/locations/%s/jobs/%s/streams/_default",
jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
location,
jobId.getJob());
}

ReadRowsRequest readRowsRequest =
ReadRowsRequest.newBuilder()
Expand Down Expand Up @@ -2705,8 +2733,7 @@ && getOptions().getOpenTelemetryTracer() != null) {
}

if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
throw new IllegalArgumentException(
"Arrow results format is only supported for fast query path execution (e.g. no destination table, no custom clustering, etc.).");
return queryFallbackArrow(jobId, configuration, options);
}

return create(JobInfo.of(jobId, configuration), options);
Expand Down Expand Up @@ -2986,6 +3013,172 @@ private ArrowQueryResult createArrowQueryResultFromTable(
return ArrowQueryResultImpl.fromReadSession(readSession, jobId, client);
}

/**
* Executes a slow-path query job using {@code jobs.insert}, awaits its completion, and streams
* the result rows via the BigQuery Storage Read API in Arrow format, wrapping the decoded rows in
* a {@link TableResult}.
*
* @param jobId the job ID, or {@code null}
* @param configuration the query job configuration
* @param options query job options
* @return a {@link TableResult} containing the decoded rows and job execution metadata
* @throws InterruptedException if interrupted while awaiting job completion
* @throws BigQueryException if job execution or ReadSession creation fails
*/
private TableResult queryFallbackArrow(
JobId jobId, QueryJobConfiguration configuration, JobOption... options)
throws InterruptedException {
Job job = create(JobInfo.of(jobId, configuration), options);
Job completedJob = job.waitFor();

if (completedJob == null) {
throw new BigQueryException(0, "Job no longer exists or could not be retrieved.");
}

if (completedJob.getStatus().getError() != null) {
throw new BigQueryException(Collections.singletonList(completedJob.getStatus().getError()));
}

TableId destinationTable = null;
if (completedJob.getConfiguration() instanceof QueryJobConfiguration) {
destinationTable =
((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable();
}
if (destinationTable == null) {
destinationTable = configuration.getDestinationTable();
}
if (destinationTable == null) {
throw new BigQueryException(0, "Unable to resolve destination table for fallback query");
}

JobStatistics.QueryStatistics stats =
completedJob.getStatistics() instanceof JobStatistics.QueryStatistics
? (JobStatistics.QueryStatistics) completedJob.getStatistics()
: null;

StatementType statementType = stats != null ? stats.getStatementType() : null;
Long totalBytesBilled = stats != null ? stats.getTotalBytesBilled() : null;
Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null;
Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null;
Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null;
SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null;

String destProject =
destinationTable.getProject() != null
? destinationTable.getProject()
: (completedJob.getJobId() != null && completedJob.getJobId().getProject() != null
? completedJob.getJobId().getProject()
: getOptions().getProjectId());
String parent = String.format("projects/%s", destProject);
String srcTable =
String.format(
"projects/%s/datasets/%s/tables/%s",
destProject, destinationTable.getDataset(), destinationTable.getTable());

BigQueryReadClient client = getBigQueryReadClient();
CreateReadSessionRequest request =
CreateReadSessionRequest.newBuilder()
.setParent(parent)
.setReadSession(
ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW))
.setMaxStreamCount(1)
.build();
ReadSession readSession;
try {
readSession = client.createReadSession(request);
} catch (Exception e) {
throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e);
}

org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null;
byte[] arrowSchemaBytes = null;
if (readSession.hasArrowSchema()) {
arrowSchemaBytes = readSession.getArrowSchema().getSerializedSchema().toByteArray();
try {
arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
} catch (IOException e) {
throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e);
}
}
Schema schema =
arrowSchemaPojo != null
? ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo)
: (stats != null ? stats.getSchema() : null);

String streamName =
readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null;

if (streamName == null) {
return TableResult.newBuilder()
.setSchema(schema)
.setTotalRows(numDmlAffectedRows != null ? numDmlAffectedRows : 0L)
.setPageNoSchema(
new PageImpl<>(
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),

@gemini-code-assist gemini-code-assist Bot Sep 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of passing null as the first argument (TableId) to TableDataPageFetcher, pass destinationTable. Since destinationTable is guaranteed to be non-null at this point, passing it ensures that the fetcher is correctly initialized and avoids potential NullPointerExceptions if getNextPage() is ever invoked.

Suggested change
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
new TableDataPageFetcher(destinationTable, schema, getOptions(), null, optionMap(options)),

null,
ImmutableList.of()))
.setJobId(completedJob.getJobId())
.setRowsInPage(0L)
.setStatementType(statementType)
.setTotalBytesBilled(totalBytesBilled)
.setTotalBytesProcessed(totalBytesProcessed)
.setTotalSlotMs(totalSlotMs)
.setNumDmlAffectedRows(numDmlAffectedRows)
.setSessionInfo(sessionInfo)
.build();
}

ArrowQueryPageFetcher pageFetcher =
new ArrowQueryPageFetcher(
completedJob.getJobId(),
streamName,
schema,
arrowSchemaBytes,
arrowSchemaPojo,
getOptions(),
0L,
configuration.getMaxResults(),
optionMap(options));

Page<FieldValueList> firstPage = pageFetcher.getNextPage();
List<FieldValueList> firstPageRows =
firstPage != null ? ImmutableList.copyOf(firstPage.getValues()) : ImmutableList.of();
long rowsInPage = (long) firstPageRows.size();
Comment on lines +3144 to +3146

@gemini-code-assist gemini-code-assist Bot Sep 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Copying the entire first page of results into a new ImmutableList via ImmutableList.copyOf(firstPage.getValues()) just to get its size is inefficient, especially for large page sizes (e.g., up to 10,000 rows). Instead, we can check if the values are already a Collection to get the size in $O(1)$ time, or use Iterables.size() to avoid copying.

    long rowsInPage = 0;
    if (firstPage != null) {
      Iterable<FieldValueList> values = firstPage.getValues();
      rowsInPage = values instanceof java.util.Collection
          ? ((java.util.Collection<?>) values).size()
          : com.google.common.collect.Iterables.size(values);
    }


Table destTable = null;
try {
destTable = getTable(destinationTable);
} catch (Exception e) {
// Non-fatal table lookup failure
}
long totalRows =
numDmlAffectedRows != null
? numDmlAffectedRows
: (destTable != null && destTable.getNumRows() != null
? destTable.getNumRows().longValue()
: rowsInPage);

return TableResult.newBuilder()
.setSchema(schema)
.setTotalRows(totalRows)
.setPageNoSchema(
firstPage != null
? firstPage
: new PageImpl<>(
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),

@gemini-code-assist gemini-code-assist Bot Sep 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of passing null as the first argument (TableId) to TableDataPageFetcher, pass destinationTable. Since destinationTable is guaranteed to be non-null at this point, passing it ensures that the fetcher is correctly initialized and avoids potential NullPointerExceptions if getNextPage() is ever invoked.

Suggested change
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
new TableDataPageFetcher(destinationTable, schema, getOptions(), null, optionMap(options)),

null,
ImmutableList.of()))
.setJobId(completedJob.getJobId())
.setRowsInPage(rowsInPage)
.setStatementType(statementType)
.setTotalBytesBilled(totalBytesBilled)
.setTotalBytesProcessed(totalBytesProcessed)
.setTotalSlotMs(totalSlotMs)
.setNumDmlAffectedRows(numDmlAffectedRows)
.setSessionInfo(sessionInfo)
.build();
}

@Override
public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) {
Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
Expand Down
Loading
Loading