Skip to content

Commit 5304fbb

Browse files
pongadpongad
authored
bigquery: use tabledata.list for query (#2642)
This sped up query benchmark from 470s to 360s. This breaks "schema injection" however, so users cannot call `row.get(String)`. This will be fixed in a follow up. Updates #2591.
1 parent 7b15777 commit 5304fbb

5 files changed

Lines changed: 52 additions & 29 deletions

File tree

‎google-cloud-bigquery/src/benchmark/java/com/google/cloud/bigquery/benchmark/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ This directory contains benchmarks for BigQuery client.
33

44
## Usage
55
From the `google-cloud-bigquery` directory, run
6-
`mvn compile exec:java -Dexec.mainClass=com.google.cloud.bigquery.benchmark.Benchmark -Dexec.args="src/main/java/com/google/cloud/bigquery/benchmark/queries.json"`
6+
`mvn compile exec:java -Dexec.mainClass=com.google.cloud.bigquery.benchmark.Benchmark -Dexec.args="src/benchmark/java/com/google/cloud/bigquery/benchmark/queries.json"`
77

88
BigQuery service caches requests so the benchmark should be run
99
at least twice, disregarding the first result.

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.base.Preconditions.checkNotNull;
2020

21+
import com.google.api.gax.paging.Page;
2122
import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
2223
import com.google.api.gax.retrying.RetrySettings;
2324
import com.google.api.gax.retrying.TimedAttemptSettings;
@@ -262,7 +263,7 @@ public Job waitFor(RetryOption... waitOptions) throws InterruptedException {
262263
*
263264
* @throws BigQueryException upon failure
264265
*/
265-
public QueryResult getQueryResults(QueryResultsOption... options) {
266+
public Page<FieldValueList> getQueryResults(QueryResultsOption... options) {
266267
if (getConfiguration().getType() != Type.QUERY) {
267268
throw new UnsupportedOperationException(
268269
"Getting query results is supported only for " + Type.QUERY + " jobs");
@@ -274,11 +275,13 @@ public QueryResult getQueryResults(QueryResultsOption... options) {
274275
throw new BigQueryException(
275276
BigQueryException.UNKNOWN_CODE, "job completed with error", getStatus().getError());
276277
}
277-
QueryResponse response = bigquery.getQueryResults(getJobId(), options);
278-
return response.getResult();
278+
279+
TableId table = ((QueryJobConfiguration) getConfiguration()).getDestinationTable();
280+
// TODO(pongad): return QueryResult so we can inject schema.
281+
return bigquery.listTableData(table);
279282
}
280283

281-
QueryResponse waitForQueryResults(
284+
private QueryResponse waitForQueryResults(
282285
RetrySettings waitSettings, final QueryResultsOption... resultsOptions)
283286
throws InterruptedException {
284287
if (getConfiguration().getType() != Type.QUERY) {

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

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1117,6 +1117,7 @@ public void testQueryRequestCompleted() throws InterruptedException {
11171117
.setJobReference(queryJob.toPb())
11181118
.setId(JOB)
11191119
.setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"));
1120+
jobResponsePb.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb());
11201121
GetQueryResultsResponse responsePb =
11211122
new GetQueryResultsResponse()
11221123
.setJobReference(queryJob.toPb())
@@ -1144,21 +1145,24 @@ public void testQueryRequestCompleted() throws InterruptedException {
11441145
PROJECT, JOB, Collections.<BigQueryRpc.Option, Object>emptyMap()))
11451146
.andReturn(jobResponsePb);
11461147
EasyMock.expect(
1147-
bigqueryRpcMock.getQueryResults(PROJECT, JOB, optionMap)).andReturn(responsePb);
1148+
bigqueryRpcMock.listTableData(
1149+
PROJECT, DATASET, TABLE, Collections.<BigQueryRpc.Option, Object>emptyMap()))
1150+
.andReturn(Tuple.<String, Iterable<TableRow>>of("", ImmutableList.of(TABLE_ROW)));
11481151

11491152
EasyMock.replay(bigqueryRpcMock);
11501153
bigquery = options.getService();
1151-
QueryResult result =
1154+
// TODO(pongad) revert back to QueryResult
1155+
Page<FieldValueList> result =
11521156
bigquery
11531157
.query(QUERY_JOB_CONFIGURATION_FOR_QUERY, queryJob)
11541158
.getQueryResults(QueryResultsOption.pageSize(42L));
1155-
assertEquals(null, result.getSchema());
1156-
assertEquals(1L, result.getTotalRows());
1159+
// assertEquals(null, result.getSchema());
1160+
// assertEquals(1L, result.getTotalRows());
11571161
for (FieldValueList row : result.getValues()) {
11581162
assertEquals(false, row.get(0).getBooleanValue());
11591163
assertEquals(1L, row.get(1).getLongValue());
11601164
}
1161-
assertEquals(CURSOR, result.getNextPageToken());
1165+
// assertEquals(CURSOR, result.getNextPageToken());
11621166
}
11631167

11641168
@Test
@@ -1169,6 +1173,7 @@ public void testQueryRequestCompletedOnSecondAttempt() throws InterruptedExcepti
11691173
.setConfiguration(QUERY_JOB_CONFIGURATION_FOR_QUERY.toPb())
11701174
.setJobReference(queryJob.toPb())
11711175
.setId(JOB);
1176+
jobResponsePb1.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb());
11721177

11731178
com.google.api.services.bigquery.model.Job jobResponsePb2 =
11741179
jobResponsePb1
@@ -1212,21 +1217,24 @@ public void testQueryRequestCompletedOnSecondAttempt() throws InterruptedExcepti
12121217
PROJECT, JOB, Collections.<BigQueryRpc.Option, Object>emptyMap()))
12131218
.andReturn(jobResponsePb2);
12141219
EasyMock.expect(
1215-
bigqueryRpcMock.getQueryResults(PROJECT, JOB, optionMap)).andReturn(responsePb2);
1220+
bigqueryRpcMock.listTableData(
1221+
PROJECT, DATASET, TABLE, Collections.<BigQueryRpc.Option, Object>emptyMap()))
1222+
.andReturn(Tuple.<String, Iterable<TableRow>>of("", ImmutableList.of(TABLE_ROW)));
12161223

12171224
EasyMock.replay(bigqueryRpcMock);
12181225
bigquery = options.getService();
1219-
QueryResult result =
1226+
// TODO(pongad) revert back to QueryResult
1227+
Page<FieldValueList> result =
12201228
bigquery
12211229
.query(QUERY_JOB_CONFIGURATION_FOR_QUERY, queryJob)
12221230
.getQueryResults(QueryResultsOption.pageSize(42L));
1223-
assertEquals(null, result.getSchema());
1224-
assertEquals(1L, result.getTotalRows());
1231+
// assertEquals(null, result.getSchema());
1232+
// assertEquals(1L, result.getTotalRows());
12251233
for (FieldValueList row : result.getValues()) {
12261234
assertEquals(false, row.get(0).getBooleanValue());
12271235
assertEquals(1L, row.get(1).getLongValue());
12281236
}
1229-
assertEquals(CURSOR, result.getNextPageToken());
1237+
// assertEquals(CURSOR, result.getNextPageToken());
12301238
}
12311239

12321240
@Test

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,8 @@ public void testWaitFor() throws InterruptedException {
222222

223223
@Test
224224
public void testWaitForAndGetQueryResults() throws InterruptedException {
225-
QueryJobConfiguration jobConfig = QueryJobConfiguration.of("SELECT 1");
225+
QueryJobConfiguration jobConfig =
226+
QueryJobConfiguration.newBuilder("SELECT 1").setDestinationTable(TABLE_ID1).build();
226227
QueryStatistics jobStatistics =
227228
QueryStatistics.newBuilder()
228229
.setCreationTimestamp(1L)
@@ -252,7 +253,7 @@ public void testWaitForAndGetQueryResults() throws InterruptedException {
252253

253254
expect(bigquery.getQueryResults(jobInfo.getJobId(), Job.DEFAULT_QUERY_WAIT_OPTIONS)).andReturn(completedQuery);
254255
expect(bigquery.getJob(JOB_INFO.getJobId())).andReturn(completedJob);
255-
expect(bigquery.getQueryResults(jobInfo.getJobId())).andReturn(completedQuery);
256+
expect(bigquery.listTableData(TABLE_ID1)).andReturn(result);
256257

257258
replay(status, bigquery, mockOptions);
258259
initializeJob(jobInfo);

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

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,6 @@
5858
import com.google.cloud.bigquery.LoadJobConfiguration;
5959
import com.google.cloud.bigquery.QueryJobConfiguration;
6060
import com.google.cloud.bigquery.QueryParameterValue;
61-
import com.google.cloud.bigquery.QueryResult;
6261
import com.google.cloud.bigquery.Schema;
6362
import com.google.cloud.bigquery.StandardTableDefinition;
6463
import com.google.cloud.bigquery.Table;
@@ -427,7 +426,8 @@ public void testCreateExternalTable() throws InterruptedException {
427426
.setDefaultDataset(DatasetId.of(DATASET))
428427
.setUseLegacySql(true)
429428
.build();
430-
QueryResult result = bigquery.query(config).getQueryResults();
429+
// TODO(pongad) revert back to QueryResult
430+
Page<FieldValueList> result = bigquery.query(config).getQueryResults();
431431
long integerValue = 0;
432432
int rowCount = 0;
433433
for (FieldValueList row : result.getValues()) {
@@ -486,7 +486,9 @@ public void testCreateViewTable() throws InterruptedException {
486486
.setDefaultDataset(DatasetId.of(DATASET))
487487
.setUseLegacySql(true)
488488
.build();
489-
QueryResult result = bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
489+
// TODO(pongad) revert back to QueryResult
490+
Page<FieldValueList> result =
491+
bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
490492
int rowCount = 0;
491493
for (FieldValueList row : result.getValues()) {
492494
FieldValue timestampCell = row.get(0);
@@ -773,8 +775,10 @@ public void testQuery() throws InterruptedException {
773775
QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query)
774776
.setDefaultDataset(DatasetId.of(DATASET))
775777
.build();
776-
QueryResult result = bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
777-
assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
778+
// TODO(pongad) revert back to QueryResult
779+
Page<FieldValueList> result =
780+
bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
781+
// assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
778782
int rowCount = 0;
779783
for (FieldValueList row : result.getValues()) {
780784
FieldValue timestampCell = row.get(0);
@@ -824,8 +828,10 @@ public void testPositionalQueryParameters() throws InterruptedException {
824828
.addPositionalParameter(int64Parameter)
825829
.addPositionalParameter(float64Parameter)
826830
.build();
827-
QueryResult result = bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
828-
assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
831+
// TODO(pongad) revert back to QueryResult
832+
Page<FieldValueList> result =
833+
bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
834+
// assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
829835
assertEquals(2, Iterables.size(result.getValues()));
830836
}
831837

@@ -844,8 +850,10 @@ public void testNamedQueryParameters() throws InterruptedException {
844850
.addNamedParameter("stringParam", stringParameter)
845851
.addNamedParameter("integerList", intArrayParameter)
846852
.build();
847-
QueryResult result = bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
848-
assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
853+
// TODO(pongad) revert back to QueryResult
854+
Page<FieldValueList> result =
855+
bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
856+
// assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
849857
assertEquals(2, Iterables.size(result.getValues()));
850858
}
851859

@@ -858,7 +866,9 @@ public void testBytesParameter() throws Exception {
858866
.setUseLegacySql(false)
859867
.addNamedParameter("p", bytesParameter)
860868
.build();
861-
QueryResult result = bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
869+
// TODO(pongad) revert back to QueryResult
870+
Page<FieldValueList> result =
871+
bigquery.query(config).getQueryResults(QueryResultsOption.pageSize(1000L));
862872
int rowCount = 0;
863873
for (FieldValueList row : result.getValues()) {
864874
rowCount++;
@@ -1021,8 +1031,9 @@ public void testQueryJob() throws InterruptedException, TimeoutException {
10211031
remoteJob = remoteJob.waitFor();
10221032
assertNull(remoteJob.getStatus().getError());
10231033

1024-
QueryResult result = remoteJob.getQueryResults();
1025-
assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
1034+
// TODO(pongad) revert back to QueryResult
1035+
Page<FieldValueList> result = remoteJob.getQueryResults();
1036+
// assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
10261037
int rowCount = 0;
10271038
for (FieldValueList row : result.getValues()) {
10281039
FieldValue timestampCell = row.get(0);

0 commit comments

Comments
 (0)