From 4d1e4eee92cef889f9112a12564a81e54e8b934c Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 15:44:20 -0400 Subject: [PATCH 1/2] test(bigquery): add ITBigQueryTest integration tests for queryArrow --- .../cloud/bigquery/it/ITBigQueryTest.java | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 9744eebc2d81..f90a72bd9795 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -45,6 +45,7 @@ import com.google.cloud.bigquery.Acl.DatasetAclEntity; import com.google.cloud.bigquery.Acl.Expr; import com.google.cloud.bigquery.Acl.User; +import com.google.cloud.bigquery.ArrowQueryResult; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.DatasetField; import com.google.cloud.bigquery.BigQuery.DatasetListOption; @@ -118,6 +119,7 @@ import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; import com.google.cloud.bigquery.QueryJobConfiguration.Priority; import com.google.cloud.bigquery.QueryParameterValue; +import com.google.cloud.bigquery.QueryResultsFormat; import com.google.cloud.bigquery.Range; import com.google.cloud.bigquery.RangePartitioning; import com.google.cloud.bigquery.Routine; @@ -203,6 +205,7 @@ import java.util.concurrent.TimeoutException; import java.util.logging.Level; import java.util.logging.Logger; +import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -7497,6 +7500,50 @@ void testQueryWithTimeout() throws InterruptedException { assertTrue(millis < 1_000_000 * 2); } + @Test + void testQueryResultsFormatArrow() throws InterruptedException { + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + BigQuery bigQuery = bigqueryHelper.getOptions().getService(); + String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts"; + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query) + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + assertNotNull(result); + int batchCount = 0; + long totalRows = 0; + for (VectorSchemaRoot root : result) { + batchCount++; + totalRows += root.getRowCount(); + assertEquals(1, root.getRowCount()); + } + assertTrue(batchCount > 0); + assertEquals(1, totalRows); + } + } + + @Test + void testQueryResultsFormatArrowMultiPage() throws InterruptedException { + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + BigQuery bigQuery = bigqueryHelper.getOptions().getService(); + String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x"; + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query) + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + assertNotNull(result); + long totalRows = 0; + for (VectorSchemaRoot root : result) { + totalRows += root.getRowCount(); + } + assertEquals(15000, totalRows); + } + } + @Test void testUniverseDomainWithInvalidUniverseDomain() { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); From 63af688f5cacb2e7e5e5646aad2c24436649b2c5 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 17:15:37 -0400 Subject: [PATCH 2/2] test(bigquery): use shared BigQuery instance in Arrow integration tests --- .../java/com/google/cloud/bigquery/it/ITBigQueryTest.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index f90a72bd9795..2a576a47d603 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -7502,15 +7502,13 @@ void testQueryWithTimeout() throws InterruptedException { @Test void testQueryResultsFormatArrow() throws InterruptedException { - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); - BigQuery bigQuery = bigqueryHelper.getOptions().getService(); String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts"; QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query) .setQueryResultsFormat(QueryResultsFormat.ARROW) .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) .build(); - try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + try (ArrowQueryResult result = bigquery.queryArrow(config)) { assertNotNull(result); int batchCount = 0; long totalRows = 0; @@ -7526,15 +7524,13 @@ void testQueryResultsFormatArrow() throws InterruptedException { @Test void testQueryResultsFormatArrowMultiPage() throws InterruptedException { - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); - BigQuery bigQuery = bigqueryHelper.getOptions().getService(); String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x"; QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query) .setQueryResultsFormat(QueryResultsFormat.ARROW) .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) .build(); - try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + try (ArrowQueryResult result = bigquery.queryArrow(config)) { assertNotNull(result); long totalRows = 0; for (VectorSchemaRoot root : result) {