Skip to content

Commit 4505e95

Browse files
authored
test(bigquery): add ITBigQueryTest integration tests for queryArrow (#14403)
This PR adds integration tests verifying the `queryArrow` client API against the live BigQuery service. It tests single-page and multi-page Arrow stream consumption, validating row counts and VectorSchemaRoot iteration over live queries.
1 parent 21e6dc8 commit 4505e95

2 files changed

Lines changed: 55 additions & 1 deletion

File tree

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,10 @@
2424
import com.google.api.core.BetaApi;
2525
import com.google.api.core.InternalApi;
2626
import com.google.api.gax.core.FixedCredentialsProvider;
27+
import com.google.api.gax.core.NoCredentialsProvider;
2728
import com.google.api.gax.paging.Page;
2829
import com.google.api.gax.rpc.HeaderProvider;
30+
import com.google.api.gax.rpc.NoHeaderProvider;
2931
import com.google.api.services.bigquery.model.ErrorProto;
3032
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
3133
import com.google.api.services.bigquery.model.ProjectList;
@@ -371,8 +373,10 @@ private static void configureReadSettings(
371373
if (options.getCredentials() != null) {
372374
settingsBuilder.setCredentialsProvider(
373375
FixedCredentialsProvider.create(options.getCredentials()));
376+
} else {
377+
settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create());
374378
}
375-
HeaderProvider headerProvider = options.getMergedHeaderProvider(null);
379+
HeaderProvider headerProvider = options.getMergedHeaderProvider(new NoHeaderProvider());
376380
if (headerProvider != null) {
377381
settingsBuilder.setHeaderProvider(headerProvider);
378382
}

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import com.google.cloud.bigquery.Acl.DatasetAclEntity;
4646
import com.google.cloud.bigquery.Acl.Expr;
4747
import com.google.cloud.bigquery.Acl.User;
48+
import com.google.cloud.bigquery.ArrowQueryResult;
4849
import com.google.cloud.bigquery.BigQuery;
4950
import com.google.cloud.bigquery.BigQuery.DatasetField;
5051
import com.google.cloud.bigquery.BigQuery.DatasetListOption;
@@ -118,6 +119,7 @@
118119
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
119120
import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
120121
import com.google.cloud.bigquery.QueryParameterValue;
122+
import com.google.cloud.bigquery.QueryResultsFormat;
121123
import com.google.cloud.bigquery.Range;
122124
import com.google.cloud.bigquery.RangePartitioning;
123125
import com.google.cloud.bigquery.Routine;
@@ -203,6 +205,7 @@
203205
import java.util.concurrent.TimeoutException;
204206
import java.util.logging.Level;
205207
import java.util.logging.Logger;
208+
import org.apache.arrow.vector.VectorSchemaRoot;
206209
import org.junit.jupiter.api.AfterAll;
207210
import org.junit.jupiter.api.BeforeAll;
208211
import org.junit.jupiter.api.Test;
@@ -7499,6 +7502,53 @@ void testQueryWithTimeout() throws InterruptedException {
74997502
assertTrue(millis < 1_000_000 * 2);
75007503
}
75017504

7505+
@Test
7506+
void testQueryResultsFormatArrow() throws InterruptedException {
7507+
String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts";
7508+
QueryJobConfiguration config =
7509+
QueryJobConfiguration.newBuilder(query)
7510+
.setQueryResultsFormat(QueryResultsFormat.ARROW)
7511+
.setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
7512+
.build();
7513+
try (ArrowQueryResult result = bigquery.queryArrow(config)) {
7514+
assertNotNull(result);
7515+
int batchCount = 0;
7516+
long totalRows = 0;
7517+
for (VectorSchemaRoot root : result) {
7518+
batchCount++;
7519+
totalRows += root.getRowCount();
7520+
assertEquals(1, root.getRowCount());
7521+
}
7522+
assertEquals(1, batchCount);
7523+
assertEquals(1, totalRows);
7524+
}
7525+
}
7526+
7527+
@Test
7528+
void testQueryResultsFormatArrowMultiPage() throws InterruptedException {
7529+
// Under fast-query execution, the initial REST response defaults to a 10 MB payload limit.
7530+
// Setting maxResults limits the initial page to 5,000 rows, forcing the remaining 10,000 rows
7531+
// to stream across multiple batches via the BigQuery Storage Read API.
7532+
String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x";
7533+
QueryJobConfiguration config =
7534+
QueryJobConfiguration.newBuilder(query)
7535+
.setQueryResultsFormat(QueryResultsFormat.ARROW)
7536+
.setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
7537+
.setMaxResults(5000L)
7538+
.build();
7539+
try (ArrowQueryResult result = bigquery.queryArrow(config)) {
7540+
assertNotNull(result);
7541+
int batchCount = 0;
7542+
long totalRows = 0;
7543+
for (VectorSchemaRoot root : result) {
7544+
batchCount++;
7545+
totalRows += root.getRowCount();
7546+
}
7547+
assertTrue(batchCount > 1);
7548+
assertEquals(15000, totalRows);
7549+
}
7550+
}
7551+
75027552
@Test
75037553
void testUniverseDomainWithInvalidUniverseDomain() {
75047554
RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();

0 commit comments

Comments
 (0)