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 b12fb598b2e2..c5c8ae086a8a 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 @@ -303,6 +303,7 @@ static class ArrowQueryPageFetcher implements NextPageFetcher { private static final long DEFAULT_PAGE_SIZE = 10000L; private final JobId jobId; + private final String customStreamName; private final Schema schema; private final byte[] arrowSchemaBytes; private final BigQueryOptions serviceOptions; @@ -326,7 +327,30 @@ static class ArrowQueryPageFetcher implements NextPageFetcher { long initialRowOffset, Long maxResults, Map optionsMap) { + this( + jobId, + null, + schema, + arrowSchemaBytes, + arrowSchemaPojo, + serviceOptions, + initialRowOffset, + maxResults, + optionsMap); + } + + ArrowQueryPageFetcher( + JobId jobId, + String customStreamName, + Schema schema, + byte[] arrowSchemaBytes, + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo, + BigQueryOptions serviceOptions, + long initialRowOffset, + Long maxResults, + Map optionsMap) { this.jobId = jobId; + this.customStreamName = customStreamName; this.schema = schema; this.arrowSchemaBytes = arrowSchemaBytes; this.arrowSchemaPojo = arrowSchemaPojo; @@ -355,13 +379,6 @@ public Page getNextPage() { List rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L)); try { - String location = - jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation(); - if (location == null) { - throw new BigQueryException( - 0, "Location must be specified to read Arrow rows from storage stream"); - } - if (bqReadClient == null) { BigQuery service = serviceOptions.getService(); if (service instanceof BigQueryImpl) { @@ -374,12 +391,23 @@ public Page getNextPage() { } if (streamIterator == null) { - String streamName = - String.format( - "projects/%s/locations/%s/jobs/%s/streams/_default", - jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(), - location, - jobId.getJob()); + String streamName; + if (customStreamName != null) { + streamName = customStreamName; + } else { + String location = + jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation(); + if (location == null) { + throw new BigQueryException( + 0, "Location must be specified to read Arrow rows from storage stream"); + } + streamName = + String.format( + "projects/%s/locations/%s/jobs/%s/streams/_default", + jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(), + location, + jobId.getJob()); + } ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder() @@ -2710,8 +2738,7 @@ && getOptions().getOpenTelemetryTracer() != null) { } if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) { - throw new UnsupportedOperationException( - "Arrow results format for slow query path execution is not yet supported."); + return queryFallbackArrow(jobId, configuration, options); } return create(JobInfo.of(jobId, configuration), options); @@ -2991,6 +3018,172 @@ private ArrowQueryResult createArrowQueryResultFromTable( return ArrowQueryResultImpl.fromReadSession(readSession, jobId, client); } + /** + * Executes a slow-path query job using {@code jobs.insert}, awaits its completion, and streams + * the result rows via the BigQuery Storage Read API in Arrow format, wrapping the decoded rows in + * a {@link TableResult}. + * + * @param jobId the job ID, or {@code null} + * @param configuration the query job configuration + * @param options query job options + * @return a {@link TableResult} containing the decoded rows and job execution metadata + * @throws InterruptedException if interrupted while awaiting job completion + * @throws BigQueryException if job execution or ReadSession creation fails + */ + private TableResult queryFallbackArrow( + JobId jobId, QueryJobConfiguration configuration, JobOption... options) + throws InterruptedException { + Job job = create(JobInfo.of(jobId, configuration), options); + Job completedJob = job.waitFor(); + + if (completedJob == null) { + throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); + } + + if (completedJob.getStatus().getError() != null) { + throw new BigQueryException(Collections.singletonList(completedJob.getStatus().getError())); + } + + TableId destinationTable = null; + if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { + destinationTable = + ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable(); + } + if (destinationTable == null) { + destinationTable = configuration.getDestinationTable(); + } + if (destinationTable == null) { + throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); + } + + JobStatistics.QueryStatistics stats = + completedJob.getStatistics() instanceof JobStatistics.QueryStatistics + ? (JobStatistics.QueryStatistics) completedJob.getStatistics() + : null; + + StatementType statementType = stats != null ? stats.getStatementType() : null; + Long totalBytesBilled = stats != null ? stats.getTotalBytesBilled() : null; + Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null; + Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null; + Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null; + SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null; + + String destProject = + destinationTable.getProject() != null + ? destinationTable.getProject() + : (completedJob.getJobId() != null && completedJob.getJobId().getProject() != null + ? completedJob.getJobId().getProject() + : getOptions().getProjectId()); + String parent = String.format("projects/%s", destProject); + String srcTable = + String.format( + "projects/%s/datasets/%s/tables/%s", + destProject, destinationTable.getDataset(), destinationTable.getTable()); + + BigQueryReadClient client = getBigQueryReadClient(); + CreateReadSessionRequest request = + CreateReadSessionRequest.newBuilder() + .setParent(parent) + .setReadSession( + ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) + .setMaxStreamCount(1) + .build(); + ReadSession readSession; + try { + readSession = client.createReadSession(request); + } catch (Exception e) { + throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e); + } + + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null; + byte[] arrowSchemaBytes = null; + if (readSession.hasArrowSchema()) { + arrowSchemaBytes = readSession.getArrowSchema().getSerializedSchema().toByteArray(); + try { + arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); + } + } + Schema schema = + arrowSchemaPojo != null + ? ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo) + : (stats != null ? stats.getSchema() : null); + + String streamName = + readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null; + + if (streamName == null) { + return TableResult.newBuilder() + .setSchema(schema) + .setTotalRows(numDmlAffectedRows != null ? numDmlAffectedRows : 0L) + .setPageNoSchema( + new PageImpl<>( + new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), + null, + ImmutableList.of())) + .setJobId(completedJob.getJobId()) + .setRowsInPage(0L) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) + .setSessionInfo(sessionInfo) + .build(); + } + + ArrowQueryPageFetcher pageFetcher = + new ArrowQueryPageFetcher( + completedJob.getJobId(), + streamName, + schema, + arrowSchemaBytes, + arrowSchemaPojo, + getOptions(), + 0L, + configuration.getMaxResults(), + optionMap(options)); + + Page firstPage = pageFetcher.getNextPage(); + List firstPageRows = + firstPage != null ? ImmutableList.copyOf(firstPage.getValues()) : ImmutableList.of(); + long rowsInPage = (long) firstPageRows.size(); + + Table destTable = null; + try { + destTable = getTable(destinationTable); + } catch (Exception e) { + // Non-fatal table lookup failure + } + long totalRows = + numDmlAffectedRows != null + ? numDmlAffectedRows + : (destTable != null && destTable.getNumRows() != null + ? destTable.getNumRows().longValue() + : rowsInPage); + + return TableResult.newBuilder() + .setSchema(schema) + .setTotalRows(totalRows) + .setPageNoSchema( + firstPage != null + ? firstPage + : new PageImpl<>( + new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), + null, + ImmutableList.of())) + .setJobId(completedJob.getJobId()) + .setRowsInPage(rowsInPage) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) + .setSessionInfo(sessionInfo) + .build(); + } + @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); 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 0dd790654d52..8fa33d7ec936 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 @@ -74,8 +74,11 @@ 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.CreateReadSessionRequest; import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; +import com.google.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadStream; import com.google.common.base.Function; import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; @@ -2949,19 +2952,126 @@ void testQueryArrowDefaultsToJobCreationOptional() throws IOException, Interrupt } @Test - void testQueryArrowResultsFormatUnsupportedConfiguration() { + void testQueryWithArrowFormatSlowPathFallback() throws Exception { + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.JobStatus jobStatus = + new com.google.api.services.bigquery.model.JobStatus().setState("DONE"); + + com.google.api.services.bigquery.model.Job jobResponsePb = + new com.google.api.services.bigquery.model.Job() + .setJobReference(queryJob.toPb()) + .setStatus(jobStatus) + .setConfiguration( + new com.google.api.services.bigquery.model.JobConfiguration() + .setQuery( + new JobConfigurationQuery() + .setQuery("SELECT id FROM test") + .setDestinationTable(TABLE_ID.toPb()))) + .setStatistics( + new com.google.api.services.bigquery.model.JobStatistics() + .setTotalSlotMs(50L) + .setQuery( + new com.google.api.services.bigquery.model.JobStatistics2() + .setStatementType("SELECT") + .setTotalBytesBilled(100L) + .setTotalBytesProcessed(200L))); + + when(bigqueryRpcMock.createSkipExceptionTranslation( + any(com.google.api.services.bigquery.model.Job.class), any())) + .thenReturn(jobResponsePb); + when(bigqueryRpcMock.getJobSkipExceptionTranslation(eq(PROJECT), eq(JOB), any(), any())) + .thenReturn(jobResponsePb); + when(bigqueryRpcMock.getQueryResultsSkipExceptionTranslation( + eq(PROJECT), eq(JOB), any(), any())) + .thenReturn( + new GetQueryResultsResponse().setJobComplete(true).setJobReference(queryJob.toPb())); + + 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(); + } + + byte[] batchBytes; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(1); + idVector.set(0, 42L); + 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); + batchBytes = out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + + @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); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + ReadSession readSession = + ReadSession.newBuilder() + .setName("projects/" + PROJECT + "/locations/" + LOCATION + "/sessions/session-1") + .setArrowSchema( + com.google.cloud.bigquery.storage.v1.ArrowSchema.newBuilder() + .setSerializedSchema(ByteString.copyFrom(schemaBytes))) + .addStreams(ReadStream.newBuilder().setName("stream-1")) + .build(); + + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockReadClient.createReadSession(any(CreateReadSessionRequest.class))) + .thenReturn(readSession); + when(mockReadClient.readRowsCallable()).thenReturn(mockCallable); + + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + QueryJobConfiguration config = - QueryJobConfiguration.newBuilder("SELECT 1") + QueryJobConfiguration.newBuilder("SELECT id FROM test") .setQueryResultsFormat(QueryResultsFormat.ARROW) - .setDestinationTable(TableId.of("dataset", "table")) + .setDestinationTable(TABLE_ID) .build(); - bigquery = options.getService(); - UnsupportedOperationException exception = - assertThrows(UnsupportedOperationException.class, () -> bigquery.query(config)); - assertTrue( - exception - .getMessage() - .contains("Arrow results format for slow query path execution is not yet supported.")); + + TableResult result = bigquery.query(config); + assertNotNull(result); + assertEquals(1, Iterables.size(result.getValues())); + assertEquals(queryJob, result.getJobId()); + assertEquals("42", result.getValues().iterator().next().get(0).getStringValue()); + assertEquals(StatementType.SELECT, result.getStatementType()); + assertEquals(100L, result.getTotalBytesBilled()); + assertEquals(200L, result.getTotalBytesProcessed()); + assertEquals(50L, result.getTotalSlotMs()); + + verify(bigqueryRpcMock) + .createSkipExceptionTranslation( + any(com.google.api.services.bigquery.model.Job.class), any()); + verify(mockReadClient).createReadSession(any(CreateReadSessionRequest.class)); + verify(mockCallable).call(any(ReadRowsRequest.class)); } @Test