diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index b0b5756c6866..07b8bef1c9d3 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -68,6 +68,7 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.net.HostAndPort; +import com.google.common.primitives.Longs; import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; @@ -2114,7 +2115,7 @@ public TableDataList call() throws IOException { } } - private static Iterable transformTableData( + private static List transformTableData( Iterable tableDataPb, final Schema schema, boolean useInt64Timestamps) { return ImmutableList.copyOf( Iterables.transform( @@ -2394,7 +2395,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options)), cursor, transformTableData( - results.getRows(), + results.getRows() != null ? results.getRows() : ImmutableList.of(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()))) .setJobId(jobId) @@ -2418,7 +2419,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), null, transformTableData( - results.getRows(), + results.getRows() != null ? results.getRows() : ImmutableList.of(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()))) // Return the JobID of the successful job @@ -2436,6 +2437,198 @@ public com.google.api.services.bigquery.model.QueryResponse call() .build(); } + /** + * Executes a fast-path query RPC request expecting Arrow-formatted wire response. + * + *

Deserializes the returned Arrow IPC schema and record batch into standard {@link + * TableResult} row-based representations, configuring {@link ArrowQueryPageFetcher} for + * subsequent pages when pagination tokens are returned. + * + * @param projectId project ID in which to execute the query + * @param content query request content PB + * @param options job options + * @return either a {@link Job} if incomplete/fallback is required, or {@link TableResult} + * @throws InterruptedException if interrupted while awaiting RPC execution + */ + private Object queryRpcArrow( + final String projectId, final QueryRequest content, JobOption... options) + throws InterruptedException { + com.google.api.services.bigquery.model.QueryResponse results; + Span queryRpc = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + queryRpc = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.queryRpc") + .setAttribute("bq.query.project_id", projectId) + .setAllAttributes(otelAttributesFromQueryRequest(content)) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryRpcScope = queryRpc != null ? queryRpc.makeCurrent() : null) { + results = + BigQueryRetryHelper.runWithRetries( + () -> bigQueryRpc.queryRpcSkipExceptionTranslation(projectId, content), + getOptions().getRetrySettings(), + getOptions().getResultRetryAlgorithm(), + getOptions().getClock(), + DEFAULT_RETRY_CONFIG, + getOptions().isOpenTelemetryTracingEnabled(), + getOptions().getOpenTelemetryTracer()); + } catch (BigQueryRetryHelper.BigQueryRetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } finally { + if (queryRpc != null) { + queryRpc.end(); + } + } + + if (results.getErrors() != null) { + List bigQueryErrors = + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); + throw new BigQueryException(bigQueryErrors); + } + + // If query is incomplete or did not return an Arrow schema, fallback to fetching the Job. + if (!Boolean.TRUE.equals(results.getJobComplete()) || results.getArrowSchema() == null) { + if (results.getJobReference() == null) { + throw new BigQueryException( + 0, "Job is incomplete or Arrow schema is missing, but no job reference was returned."); + } + JobId jobId = JobId.fromPb(results.getJobReference()); + return getJob(jobId, options); + } + + if (results.getArrowSchema().getSerializedSchema() == null) { + throw new BigQueryException(0, "Arrow schema is missing from the response"); + } + + // Deserialize Arrow IPC schema and convert to veneer BigQuery Schema. + byte[] arrowSchemaBytes = results.getArrowSchema().decodeSerializedSchema(); + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo; + try { + arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); + } + Schema schema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo); + + long numRows; + if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { + numRows = 0L; + } else if (results.getNumDmlAffectedRows() != null) { + numRows = results.getNumDmlAffectedRows(); + } else { + numRows = results.getTotalRows().longValue(); + } + + StatementType statementType = + results.getStatementType() != null + ? StatementType.valueOf(results.getStatementType()) + : null; + Long totalBytesBilled = results.getTotalBytesBilled(); + Long totalBytesProcessed = results.getTotalBytesProcessed(); + Long totalSlotMs = results.getTotalSlotMs(); + Long numDmlAffectedRows = results.getNumDmlAffectedRows(); + SessionInfo sessionInfo = + results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null; + + // Deserialize first page of rows from the Arrow record batch (if present). + List firstPageRows; + if (results.getArrowRecordBatch() == null + || results.getArrowRecordBatch().getSerializedRecordBatch() == null) { + firstPageRows = ImmutableList.of(); + } else { + try { + firstPageRows = + ArrowDeserializer.deserializeRecordBatch( + results.getArrowRecordBatch().decodeSerializedRecordBatch(), + schema, + arrowSchemaPojo); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e); + } + } + + // Enforce maxResults limit on the first page if requested. + if (content.getMaxResults() != null && firstPageRows.size() > content.getMaxResults()) { + firstPageRows = + ImmutableList.copyOf(Iterables.limit(firstPageRows, content.getMaxResults().intValue())); + } + + // Calculate row offset and determine if subsequent pages exist. + boolean hasMorePages = results.getPageToken() != null; + long initialRowOffset = 0L; + if (hasMorePages) { + Long parsedOffset = Longs.tryParse(results.getPageToken()); + initialRowOffset = parsedOffset != null ? parsedOffset : firstPageRows.size(); + if (content.getMaxResults() != null + && (initialRowOffset >= content.getMaxResults() + || firstPageRows.size() >= content.getMaxResults())) { + hasMorePages = false; + } + } + + // Multi-page results: configure ArrowQueryPageFetcher for subsequent tabledata.list calls. + if (hasMorePages) { + if (results.getJobReference() == null) { + throw new BigQueryException( + 0, "More pages exist, but no job reference was returned to fetch them."); + } + JobId jobId = JobId.fromPb(results.getJobReference()); + String cursor = results.getPageToken(); + NextPageFetcher pageFetcher = + new ArrowQueryPageFetcher( + jobId, + schema, + arrowSchemaBytes, + arrowSchemaPojo, + getOptions(), + initialRowOffset, + content.getMaxResults(), + optionMap(options)); + + return TableResult.newBuilder() + .setSchema(schema) + .setTotalRows(numRows) + .setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows)) + .setJobId(jobId) + .setQueryId(results.getQueryId()) + .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason())) + .setRowsInPage((long) firstPageRows.size()) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) + .setSessionInfo(sessionInfo) + .build(); + } + + // only 1 page of result + return TableResult.newBuilder() + .setSchema(schema) + .setTotalRows(numRows) + .setPageNoSchema( + new PageImpl<>( + new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), + null, + firstPageRows)) + .setJobId( + results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null) + .setQueryId(results.getQueryId()) + .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason())) + .setRowsInPage((long) firstPageRows.size()) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) + .setSessionInfo(sessionInfo) + .build(); + } + @Override public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options) throws InterruptedException, JobException { @@ -2452,11 +2645,6 @@ public Object queryWithTimeout( throws InterruptedException, JobException { Job.checkNotDryRun(configuration, "query"); - if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) { - throw new IllegalArgumentException( - "QueryResultsFormat.ARROW is not supported with query(). Use queryArrow() instead."); - } - // If JobCreationMode is not explicitly set, update it with default value; if (configuration.getJobCreationMode() == null) { configuration = @@ -2509,8 +2697,17 @@ && getOptions().getOpenTelemetryTracer() != null) { content.setTimeoutMs(timeoutMs); } + if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) { + return queryRpcArrow(projectId, content, options); + } return queryRpc(projectId, content, options); } + + 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 create(JobInfo.of(jobId, configuration), options); } finally { if (querySpan != null) { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index cb67ac54aa4a..9428456748e8 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -37,6 +37,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import com.google.api.client.googleapis.json.GoogleJsonError; import com.google.api.client.googleapis.json.GoogleJsonResponseException; @@ -45,6 +46,8 @@ import com.google.api.gax.paging.Page; import com.google.api.gax.retrying.ResultRetryAlgorithm; import com.google.api.gax.retrying.TimedAttemptSettings; +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; import com.google.api.services.bigquery.model.JobConfigurationQuery; @@ -70,6 +73,9 @@ import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.spi.v2.BigQueryRpc; import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc; +import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.common.base.Function; import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; @@ -77,13 +83,27 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; +import com.google.common.io.BaseEncoding; +import com.google.protobuf.ByteString; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.math.BigInteger; import java.net.ConnectException; import java.net.UnknownHostException; +import java.nio.channels.Channels; import java.util.Collections; import java.util.List; import java.util.Map; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.WriteChannel; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -2905,41 +2925,372 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException } @Test - void testQueryThrowsWhenArrowResultsFormat() { + void testQueryArrowDefaultsToJobCreationOptional() throws IOException, InterruptedException { + QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1").build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-optional-1") + .setJobComplete(true) + .setTotalRows(BigInteger.ZERO); + + ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + ArrowQueryResult result = bigquery.queryArrow(config); + assertNotNull(result); + assertEquals("q-optional-1", result.getQueryId()); + assertNull(result.getJobId()); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals("JOB_CREATION_OPTIONAL", requestPb.getJobCreationMode()); + assertEquals("ARROW", requestPb.getQueryResultsFormat()); + } + + @Test + void testQueryArrowResultsFormatUnsupportedConfiguration() { QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1") .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setDestinationTable(TableId.of("dataset", "table")) .build(); bigquery = options.getService(); IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); - assertTrue(exception.getMessage().contains("Use queryArrow() instead")); + assertTrue( + exception + .getMessage() + .contains("Arrow results format is only supported for fast query path execution")); } @Test - void testQueryArrowDefaultsToJobCreationOptional() throws IOException, InterruptedException { - QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1").build(); + void testQueryWithArrowFormatFastPath() throws IOException, InterruptedException { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of( + org.apache.arrow.vector.types.pojo.Field.nullable( + "id", new ArrowType.Int(64, true)))); + + byte[] schemaBytes; + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema); + schemaBytes = out.toByteArray(); + } + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = new com.google.api.services.bigquery.model.QueryResponse() - .setQueryId("q-optional-1") + .setQueryId("q-arrow-1") .setJobComplete(true) - .setTotalRows(java.math.BigInteger.ZERO); + .setTotalRows(BigInteger.ONE) + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) .thenReturn(queryResponsePb); + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT 1 as id") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); bigquery = options.getService(); - ArrowQueryResult result = bigquery.queryArrow(config); + TableResult result = bigquery.query(config); assertNotNull(result); - assertEquals("q-optional-1", result.getQueryId()); - assertNull(result.getJobId()); + assertEquals("q-arrow-1", result.getQueryId()); + assertNotNull(result.getSchema()); + assertEquals(1, result.getSchema().getFields().size()); + assertEquals("id", result.getSchema().getFields().get(0).getName()); QueryRequest requestPb = requestPbCapture.getValue(); - assertEquals("JOB_CREATION_OPTIONAL", requestPb.getJobCreationMode()); assertEquals("ARROW", requestPb.getQueryResultsFormat()); } + @Test + void testQueryWithArrowFormatMultiplePages() throws IOException, InterruptedException { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of( + org.apache.arrow.vector.types.pojo.Field.nullable( + "id", new ArrowType.Int(64, true)))); + + byte[] schemaBytes; + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema); + schemaBytes = out.toByteArray(); + } + + // Prepare page 2 Arrow batch for streaming + byte[] page2BatchBytes; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(1); + idVector.set(0, 2L); + idVector.setValueCount(1); + try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) { + VectorUnloader unloader = new VectorUnloader(root); + try (ArrowRecordBatch recordBatch = unloader.getRecordBatch(); + ByteArrayOutputStream out = new ByteArrayOutputStream()) { + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + MessageSerializer.serialize(channel, recordBatch); + page2BatchBytes = out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-multipage") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + // Mock BigQueryReadClient for page 2 + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); + @SuppressWarnings("unchecked") + ServerStream mockServerStream = + mock(ServerStream.class, withSettings().withoutAnnotations()); + when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockReadClient.readRowsCallable()).thenReturn(mockCallable); + + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertEquals("q-arrow-multipage", result.getQueryId()); + assertTrue(result.hasNextPage()); + assertEquals("1", result.getNextPageToken()); + + Page page2 = result.getNextPage(); + assertNotNull(page2); + List page2Rows = ImmutableList.copyOf(page2.getValues()); + assertEquals(1, page2Rows.size()); + assertEquals("2", page2Rows.get(0).get(0).getStringValue()); + + verify(mockCallable).call(any(ReadRowsRequest.class)); + } + + @Test + void testQueryWithArrowFormatMultiplePagesWithMaxResults() + throws IOException, InterruptedException { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of( + org.apache.arrow.vector.types.pojo.Field.nullable( + "id", new ArrowType.Int(64, true)))); + + byte[] schemaBytes; + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema); + schemaBytes = out.toByteArray(); + } + + // Prepare page 2 Arrow batch for streaming with 2 rows + byte[] page2BatchBytes; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(2); + idVector.set(0, 2L); + idVector.set(1, 3L); + idVector.setValueCount(2); + try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) { + VectorUnloader unloader = new VectorUnloader(root); + try (ArrowRecordBatch recordBatch = unloader.getRecordBatch(); + ByteArrayOutputStream out = new ByteArrayOutputStream()) { + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + MessageSerializer.serialize(channel, recordBatch); + page2BatchBytes = out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-multipage-maxresults") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(3L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + // Mock BigQueryReadClient for page 2 + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); + @SuppressWarnings("unchecked") + ServerStream mockServerStream = + mock(ServerStream.class, withSettings().withoutAnnotations()); + when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockReadClient.readRowsCallable()).thenReturn(mockCallable); + + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setMaxResults(2L) + .build(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertTrue(result.hasNextPage()); + + Page page2 = result.getNextPage(); + assertNotNull(page2); + List page2Rows = ImmutableList.copyOf(page2.getValues()); + // Since maxResults is 2 and initialRowOffset is 1, page2 should only contain 1 row even though + // stream returned 2 rows + assertEquals(1, page2Rows.size()); + assertEquals("2", page2Rows.get(0).get(0).getStringValue()); + // Since totalRowsReturned == maxResults, hasNextPage must be false + assertFalse(page2.hasNextPage()); + assertNull(page2.getNextPage()); + + // When maxResults is 1, initialRowOffset (1) already reaches maxResults, so hasNextPage is + // false immediately + QueryJobConfiguration configMax1 = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setMaxResults(1L) + .build(); + TableResult resultMax1 = bigquery.query(configMax1); + assertNotNull(resultMax1); + assertFalse(resultMax1.hasNextPage()); + assertNull(resultMax1.getNextPage()); + } + + @Test + void testArrowQueryPageFetcherSerialization() throws Exception { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of( + org.apache.arrow.vector.types.pojo.Field.nullable( + "id", new ArrowType.Int(64, true)))); + + byte[] schemaBytes; + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema); + schemaBytes = out.toByteArray(); + } + + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-multipage-ser") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertTrue(result.hasNextPage()); + assertEquals("1", result.getNextPageToken()); + + // Serialize and deserialize TableResult + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (java.io.ObjectOutputStream oos = new java.io.ObjectOutputStream(baos)) { + oos.writeObject(result); + } + + TableResult deserializedResult; + try (java.io.ObjectInputStream ois = + new java.io.ObjectInputStream(new ByteArrayInputStream(baos.toByteArray()))) { + deserializedResult = (TableResult) ois.readObject(); + } + + assertNotNull(deserializedResult); + assertEquals(result.getSchema(), deserializedResult.getSchema()); + assertEquals(result.getTotalRows(), deserializedResult.getTotalRows()); + assertEquals(result.getQueryId(), deserializedResult.getQueryId()); + assertEquals("1", deserializedResult.getNextPageToken()); + assertTrue(deserializedResult.hasNextPage()); + } + + @Test + void testQueryWithArrowFormatMissingSerializedSchema() throws Exception { + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-missing-schema") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setArrowSchema(new com.google.api.services.bigquery.model.ArrowSchema()); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + BigQueryException e = assertThrows(BigQueryException.class, () -> bigquery.query(config)); + assertTrue(e.getMessage().contains("Arrow schema is missing from the response")); + } + @Test void testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB);