Skip to content

Commit ee96503

Browse files
jinseopkim0jinseopkim0
committed
fix(bigquery): address review feedback on queryRpcArrow and TableResult builder
1 parent 773e182 commit ee96503

2 files changed

Lines changed: 75 additions & 74 deletions

File tree

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

Lines changed: 35 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -2403,28 +2403,40 @@ private static long getNumRows(com.google.api.services.bigquery.model.QueryRespo
24032403
}
24042404
}
24052405

2406-
private TableResult createTableResultFromJsonResponse(
2407-
com.google.api.services.bigquery.model.QueryResponse results, JobOption... options) {
2408-
long numRows = getNumRows(results);
2409-
Schema schema = results.getSchema() != null ? Schema.fromPb(results.getSchema()) : null;
2410-
2406+
private static TableResult.Builder newTableResultBuilder(
2407+
com.google.api.services.bigquery.model.QueryResponse results) {
24112408
StatementType statementType =
24122409
results.getStatementType() != null
24132410
? StatementType.valueOf(results.getStatementType())
24142411
: null;
2415-
Long totalBytesBilled = results.getTotalBytesBilled();
2416-
Long totalBytesProcessed = results.getTotalBytesProcessed();
2417-
Long totalSlotMs = results.getTotalSlotMs();
2418-
Long numDmlAffectedRows = results.getNumDmlAffectedRows();
24192412
SessionInfo sessionInfo =
24202413
results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
2414+
JobCreationReason jobCreationReason =
2415+
results.getJobCreationReason() != null
2416+
? JobCreationReason.fromPb(results.getJobCreationReason())
2417+
: null;
2418+
2419+
return TableResult.newBuilder()
2420+
.setTotalRows(getNumRows(results))
2421+
.setQueryId(results.getQueryId())
2422+
.setJobCreationReason(jobCreationReason)
2423+
.setStatementType(statementType)
2424+
.setTotalBytesBilled(results.getTotalBytesBilled())
2425+
.setTotalBytesProcessed(results.getTotalBytesProcessed())
2426+
.setTotalSlotMs(results.getTotalSlotMs())
2427+
.setNumDmlAffectedRows(results.getNumDmlAffectedRows())
2428+
.setSessionInfo(sessionInfo);
2429+
}
2430+
2431+
private TableResult createTableResultFromJsonResponse(
2432+
com.google.api.services.bigquery.model.QueryResponse results, JobOption... options) {
2433+
Schema schema = results.getSchema() != null ? Schema.fromPb(results.getSchema()) : null;
24212434

24222435
if (results.getPageToken() != null) {
24232436
JobId jobId = JobId.fromPb(results.getJobReference());
24242437
String cursor = results.getPageToken();
2425-
return TableResult.newBuilder()
2438+
return newTableResultBuilder(results)
24262439
.setSchema(schema)
2427-
.setTotalRows(numRows)
24282440
.setPageNoSchema(
24292441
new PageImpl<>(
24302442
// fetch next pages of results
@@ -2435,21 +2447,12 @@ private TableResult createTableResultFromJsonResponse(
24352447
schema,
24362448
getOptions().getDataFormatOptions().useInt64Timestamp())))
24372449
.setJobId(jobId)
2438-
.setQueryId(results.getQueryId())
2439-
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
24402450
.setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
2441-
.setStatementType(statementType)
2442-
.setTotalBytesBilled(totalBytesBilled)
2443-
.setTotalBytesProcessed(totalBytesProcessed)
2444-
.setTotalSlotMs(totalSlotMs)
2445-
.setNumDmlAffectedRows(numDmlAffectedRows)
2446-
.setSessionInfo(sessionInfo)
24472451
.build();
24482452
}
24492453
// only 1 page of result
2450-
return TableResult.newBuilder()
2454+
return newTableResultBuilder(results)
24512455
.setSchema(schema)
2452-
.setTotalRows(numRows)
24532456
.setPageNoSchema(
24542457
new PageImpl<>(
24552458
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
@@ -2461,15 +2464,7 @@ private TableResult createTableResultFromJsonResponse(
24612464
// Return the JobID of the successful job
24622465
.setJobId(
24632466
results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
2464-
.setQueryId(results.getQueryId())
2465-
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
24662467
.setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
2467-
.setStatementType(statementType)
2468-
.setTotalBytesBilled(totalBytesBilled)
2469-
.setTotalBytesProcessed(totalBytesProcessed)
2470-
.setTotalSlotMs(totalSlotMs)
2471-
.setNumDmlAffectedRows(numDmlAffectedRows)
2472-
.setSessionInfo(sessionInfo)
24732468
.build();
24742469
}
24752470

@@ -2532,13 +2527,9 @@ && getOptions().getOpenTelemetryTracer() != null) {
25322527
"Arrow results format for slow query path execution is not yet supported.");
25332528
}
25342529

2535-
// If query is complete but Arrow schema is missing or invalid, fallback to standard JSON
2536-
// TableResult if JSON schema is present, or throw an exception.
2530+
// If the query completed but the Arrow schema is missing from the response, fail fast.
25372531
if (results.getArrowSchema() == null
25382532
|| results.getArrowSchema().getSerializedSchema() == null) {
2539-
if (results.getSchema() != null) {
2540-
return createTableResultFromJsonResponse(results, options);
2541-
}
25422533
throw new BigQueryException(0, "Arrow schema is missing from the response");
25432534
}
25442535

@@ -2552,19 +2543,6 @@ && getOptions().getOpenTelemetryTracer() != null) {
25522543
}
25532544
Schema schema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo);
25542545

2555-
long numRows = getNumRows(results);
2556-
2557-
StatementType statementType =
2558-
results.getStatementType() != null
2559-
? StatementType.valueOf(results.getStatementType())
2560-
: null;
2561-
Long totalBytesBilled = results.getTotalBytesBilled();
2562-
Long totalBytesProcessed = results.getTotalBytesProcessed();
2563-
Long totalSlotMs = results.getTotalSlotMs();
2564-
Long numDmlAffectedRows = results.getNumDmlAffectedRows();
2565-
SessionInfo sessionInfo =
2566-
results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
2567-
25682546
// Deserialize first page of rows from the Arrow record batch (if present).
25692547
List<FieldValueList> firstPageRows;
25702548
if (results.getArrowRecordBatch() == null
@@ -2593,7 +2571,14 @@ && getOptions().getOpenTelemetryTracer() != null) {
25932571
long initialRowOffset = 0L;
25942572
if (hasMorePages) {
25952573
Long parsedOffset = Longs.tryParse(results.getPageToken());
2596-
initialRowOffset = parsedOffset != null ? parsedOffset : firstPageRows.size();
2574+
if (parsedOffset == null) {
2575+
throw new BigQueryException(
2576+
0,
2577+
String.format(
2578+
"Unable to parse page token '%s' as a numeric row offset for Arrow query pagination",
2579+
results.getPageToken()));
2580+
}
2581+
initialRowOffset = parsedOffset;
25972582
if (content.getMaxResults() != null
25982583
&& (initialRowOffset >= content.getMaxResults()
25992584
|| firstPageRows.size() >= content.getMaxResults())) {
@@ -2620,49 +2605,25 @@ && getOptions().getOpenTelemetryTracer() != null) {
26202605
content.getMaxResults(),
26212606
optionMap(options));
26222607

2623-
return TableResult.newBuilder()
2608+
return newTableResultBuilder(results)
26242609
.setSchema(schema)
2625-
.setTotalRows(numRows)
26262610
.setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows))
26272611
.setJobId(jobId)
2628-
.setQueryId(results.getQueryId())
2629-
.setJobCreationReason(
2630-
results.getJobCreationReason() != null
2631-
? JobCreationReason.fromPb(results.getJobCreationReason())
2632-
: null)
26332612
.setRowsInPage((long) firstPageRows.size())
2634-
.setStatementType(statementType)
2635-
.setTotalBytesBilled(totalBytesBilled)
2636-
.setTotalBytesProcessed(totalBytesProcessed)
2637-
.setTotalSlotMs(totalSlotMs)
2638-
.setNumDmlAffectedRows(numDmlAffectedRows)
2639-
.setSessionInfo(sessionInfo)
26402613
.build();
26412614
}
26422615

26432616
// only 1 page of result
2644-
return TableResult.newBuilder()
2617+
return newTableResultBuilder(results)
26452618
.setSchema(schema)
2646-
.setTotalRows(numRows)
26472619
.setPageNoSchema(
26482620
new PageImpl<>(
26492621
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
26502622
null,
26512623
firstPageRows))
26522624
.setJobId(
26532625
results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
2654-
.setQueryId(results.getQueryId())
2655-
.setJobCreationReason(
2656-
results.getJobCreationReason() != null
2657-
? JobCreationReason.fromPb(results.getJobCreationReason())
2658-
: null)
26592626
.setRowsInPage((long) firstPageRows.size())
2660-
.setStatementType(statementType)
2661-
.setTotalBytesBilled(totalBytesBilled)
2662-
.setTotalBytesProcessed(totalBytesProcessed)
2663-
.setTotalSlotMs(totalSlotMs)
2664-
.setNumDmlAffectedRows(numDmlAffectedRows)
2665-
.setSessionInfo(sessionInfo)
26662627
.build();
26672628
}
26682629

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3272,6 +3272,7 @@ void testQueryWithArrowFormatMissingSerializedSchema() throws Exception {
32723272
.setJobComplete(true)
32733273
.setJobReference(queryJob.toPb())
32743274
.setTotalRows(BigInteger.valueOf(2L))
3275+
.setSchema(TABLE_SCHEMA.toPb())
32753276
.setArrowSchema(new com.google.api.services.bigquery.model.ArrowSchema());
32763277

32773278
when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class)))
@@ -3312,6 +3313,45 @@ void testQueryWithArrowFormatIncompleteJob() throws Exception {
33123313
.contains("Arrow results format for slow query path execution is not yet supported."));
33133314
}
33143315

3316+
@Test
3317+
void testQueryWithArrowFormatInvalidPageToken() throws Exception {
3318+
org.apache.arrow.vector.types.pojo.Schema arrowSchema =
3319+
new org.apache.arrow.vector.types.pojo.Schema(
3320+
ImmutableList.of(
3321+
org.apache.arrow.vector.types.pojo.Field.nullable(
3322+
"id", new ArrowType.Int(64, true))));
3323+
3324+
byte[] schemaBytes;
3325+
try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
3326+
MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
3327+
schemaBytes = out.toByteArray();
3328+
}
3329+
3330+
JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
3331+
com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
3332+
new com.google.api.services.bigquery.model.QueryResponse()
3333+
.setQueryId("q-arrow-invalid-token")
3334+
.setJobComplete(true)
3335+
.setJobReference(queryJob.toPb())
3336+
.setTotalRows(BigInteger.valueOf(2L))
3337+
.setPageToken("invalid-non-numeric-token")
3338+
.setArrowSchema(
3339+
new com.google.api.services.bigquery.model.ArrowSchema()
3340+
.setSerializedSchema(BaseEncoding.base64().encode(schemaBytes)));
3341+
3342+
when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class)))
3343+
.thenReturn(queryResponsePb);
3344+
3345+
bigquery = options.getService();
3346+
3347+
QueryJobConfiguration config =
3348+
QueryJobConfiguration.newBuilder("SELECT id FROM test")
3349+
.setQueryResultsFormat(QueryResultsFormat.ARROW)
3350+
.build();
3351+
BigQueryException e = assertThrows(BigQueryException.class, () -> bigquery.query(config));
3352+
assertTrue(e.getMessage().contains("Unable to parse page token 'invalid-non-numeric-token'"));
3353+
}
3354+
33153355
@Test
33163356
void testGetQueryResults() throws IOException {
33173357
JobId queryJob = JobId.of(JOB);

0 commit comments

Comments
 (0)