feat(bigquery): integrate Arrow query response processing and stream pagination

Author: jinseopkim0

java-bigquery/google-cloud-bigquery/pom.xml

@@ -120,6 +120,15 @@
       <artifactId>arrow-memory-netty</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>com.google.api</groupId>
+      <artifactId>gax-grpc</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-api</artifactId>
+    </dependency>
+
     <dependency>
       <groupId>com.google.errorprone</groupId>
       <artifactId>error_prone_annotations</artifactId>

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

@@ -40,7 +40,10 @@
  *
  * @see <a href="https://cloud.google.com/bigquery/what-is-bigquery">Google Cloud BigQuery</a>
  */
-public interface BigQuery extends Service<BigQueryOptions> {
+public interface BigQuery extends Service<BigQueryOptions>, AutoCloseable {
+
+  @Override
+  default void close() {}
 
   /**
    * Fields of a BigQuery Dataset resource.

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

@@ -22,7 +22,9 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.paging.Page;
+import com.google.api.gax.rpc.ServerStream;
 import com.google.api.services.bigquery.model.ErrorProto;
 import com.google.api.services.bigquery.model.GetQueryResultsResponse;
 import com.google.api.services.bigquery.model.ProjectList;
@@ -43,6 +45,10 @@
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
 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.BigQueryReadSettings;
+import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
+import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Strings;
@@ -264,6 +270,202 @@ public Page<FieldValueList> getNextPage() {
     }
   }
 
+  private static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
+    private static final long serialVersionUID = 1L;
+
+    private static final long DEFAULT_PAGE_SIZE = 10000L;
+
+    private final JobId jobId;
+    private final Schema schema;
+    private final String arrowSchemaJson;
+    private final BigQueryOptions serviceOptions;
+    private final long maxResults;
+
+    private transient Object arrowSchemaPojo;
+    private final long totalRowsReturned;
+
+    ArrowQueryPageFetcher(
+        JobId jobId,
+        Schema schema,
+        Object arrowSchemaPojo,
+        BigQueryOptions serviceOptions,
+        long initialRowOffset,
+        Long maxResults) {
+      this.jobId = jobId;
+      this.schema = schema;
+      this.arrowSchemaJson = ArrowDeserializer.arrowSchemaToJson(arrowSchemaPojo);
+      this.arrowSchemaPojo = arrowSchemaPojo;
+      this.serviceOptions = serviceOptions;
+      this.totalRowsReturned = initialRowOffset;
+      this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE;
+    }
+
+    @Override
+    public Page<FieldValueList> getNextPage() {
+      if (totalRowsReturned >= maxResults) {
+        return null;
+      }
+
+      List<FieldValueList> rowBatch =
+          new ArrayList<>((int) Math.min(DEFAULT_PAGE_SIZE, maxResults - totalRowsReturned));
+      boolean hasMore = false;
+      BigQueryReadClient ownedClient = null;
+
+      try {
+        if (arrowSchemaPojo == null && arrowSchemaJson != null) {
+          arrowSchemaPojo = ArrowDeserializer.jsonToArrowSchema(arrowSchemaJson);
+        }
+
+        String location =
+            jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
+        if (location == null) {
+          throw new BigQueryException(
+              0, "Job location is required to read Arrow rows from storage stream");
+        }
+
+        BigQuery service = serviceOptions.getService();
+        BigQueryReadClient client;
+        if (service instanceof BigQueryImpl) {
+          client = ((BigQueryImpl) service).getBigQueryReadClient();
+        } else {
+          BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder();
+          configureReadSettings(settingsBuilder, serviceOptions);
+          ownedClient = BigQueryReadClient.create(settingsBuilder.build());
+          client = ownedClient;
+        }
+
+        String projectId =
+            jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId();
+        if (projectId == null) {
+          throw new BigQueryException(
+              0, "Project ID is required to read Arrow rows from storage stream");
+        }
+        String streamName =
+            String.format(
+                "projects/%s/locations/%s/jobs/%s/streams/_default",
+                projectId, location, jobId.getJob());
+
+        ReadRowsRequest readRowsRequest =
+            ReadRowsRequest.newBuilder()
+                .setReadStream(streamName)
+                .setOffset(totalRowsReturned)
+                .build();
+
+        ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
+        try {
+          hasMore =
+              ArrowDeserializer.loadArrowRows(
+                  stream.iterator(),
+                  arrowSchemaPojo,
+                  arrowSchemaJson,
+                  schema,
+                  rowBatch,
+                  DEFAULT_PAGE_SIZE,
+                  totalRowsReturned,
+                  maxResults);
+        } finally {
+          stream.cancel();
+        }
+      } catch (BigQueryException e) {
+        throw e;
+      } catch (Exception e) {
+        throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e);
+      } finally {
+        if (ownedClient != null) {
+          try {
+            ownedClient.close();
+          } catch (Exception e) {
+            // ignore
+          }
+        }
+      }
+
+      if (rowBatch.isEmpty()) {
+        return null;
+      }
+
+      long nextOffset = totalRowsReturned + rowBatch.size();
+      String nextPageToken = hasMore ? String.valueOf(nextOffset) : null;
+      ArrowQueryPageFetcher nextPageFetcher =
+          new ArrowQueryPageFetcher(
+              jobId, schema, arrowSchemaPojo, serviceOptions, nextOffset, maxResults);
+      return new PageImpl<>(nextPageFetcher, nextPageToken, rowBatch);
+    }
+  }
+
+  private final java.util.concurrent.locks.ReentrantLock readClientLock =
+      new java.util.concurrent.locks.ReentrantLock();
+  private transient BigQueryReadClient bqReadClient;
+
+  @Override
+  public void close() {
+    readClientLock.lock();
+    try {
+      if (bqReadClient != null) {
+        try {
+          bqReadClient.close();
+        } finally {
+          bqReadClient = null;
+        }
+      }
+    } finally {
+      readClientLock.unlock();
+    }
+  }
+
+  BigQueryReadClient getBigQueryReadClient() throws IOException {
+    readClientLock.lock();
+    try {
+      if (bqReadClient == null) {
+        BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder();
+        configureReadSettings(settingsBuilder, getOptions());
+        bqReadClient = BigQueryReadClient.create(settingsBuilder.build());
+      }
+      return bqReadClient;
+    } finally {
+      readClientLock.unlock();
+    }
+  }
+
+  private static void configureReadSettings(
+      BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options) {
+    if (options.getCredentials() != null) {
+      settingsBuilder.setCredentialsProvider(
+          FixedCredentialsProvider.create(options.getCredentials()));
+    }
+    if (options.getUniverseDomain() != null) {
+      settingsBuilder.setUniverseDomain(options.getUniverseDomain());
+    }
+    if (options.getHost() != null) {
+      String host = options.getHost();
+      String target = host;
+      if (target.contains("://")) {
+        target = java.net.URI.create(target).getAuthority();
+      }
+      com.google.common.net.HostAndPort hostAndPort =
+          com.google.common.net.HostAndPort.fromString(target);
+      String endpointHost = hostAndPort.getHost();
+      if (endpointHost.contains("bigquery.googleapis.com")) {
+        endpointHost =
+            endpointHost.replace("bigquery.googleapis.com", "bigquerystorage.googleapis.com");
+      } else if (endpointHost.contains("bigquery.private.googleapis.com")) {
+        endpointHost =
+            endpointHost.replace(
+                "bigquery.private.googleapis.com", "bigquerystorage.private.googleapis.com");
+      } else if (endpointHost.startsWith("bigquery.")) {
+        endpointHost = endpointHost.replaceFirst("^bigquery\\.", "bigquerystorage.");
+      }
+      int port = hostAndPort.getPortOrDefault(443);
+      settingsBuilder.setEndpoint(endpointHost + ":" + port);
+      if (endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1")) {
+        settingsBuilder.setTransportChannelProvider(
+            BigQueryReadSettings.defaultGrpcTransportProviderBuilder()
+                .setChannelConfigurator(io.grpc.ManagedChannelBuilder::usePlaintext)
+                .build());
+      }
+    }
+  }
+
   private final HttpBigQueryRpc bigQueryRpc;
 
   private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG =
@@ -2077,8 +2279,26 @@ public com.google.api.services.bigquery.model.QueryResponse call()
 
     long numRows;
     Schema schema;
-    if (results.getJobComplete() && results.getSchema() != null) {
-      schema = Schema.fromPb(results.getSchema());
+    boolean isArrow = false;
+    Object arrowSchemaPojo = null;
+
+    if (results.getJobComplete()) {
+      if (results.getArrowSchema() != null) {
+        isArrow = true;
+        try {
+          arrowSchemaPojo =
+              ArrowDeserializer.deserializeSchema(
+                  results.getArrowSchema().decodeSerializedSchema());
+          schema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchemaPojo);
+        } catch (IOException e) {
+          throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e);
+        }
+      } else if (results.getSchema() != null) {
+        schema = Schema.fromPb(results.getSchema());
+      } else {
+        schema = null;
+      }
+
       if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
         numRows = 0L;
       } else if (results.getNumDmlAffectedRows() != null) {
@@ -2095,45 +2315,75 @@ public com.google.api.services.bigquery.model.QueryResponse call()
       return job;
     }
 
+    List<FieldValueList> firstPageRows;
+    if (isArrow) {
+      if (results.getArrowRecordBatch() != null) {
+        try {
+          firstPageRows =
+              ArrowDeserializer.deserializeRecordBatch(
+                  results.getArrowRecordBatch().decodeSerializedRecordBatch(),
+                  schema,
+                  (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaPojo);
+        } catch (IOException e) {
+          throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e);
+        }
+      } else {
+        firstPageRows = ImmutableList.of();
+      }
+    } else {
+      firstPageRows =
+          ImmutableList.copyOf(
+              transformTableData(
+                  results.getRows(),
+                  schema,
+                  getOptions().getDataFormatOptions().useInt64Timestamp()));
+    }
+
     if (results.getPageToken() != null) {
       JobId jobId = JobId.fromPb(results.getJobReference());
       String cursor = results.getPageToken();
+
+      NextPageFetcher<FieldValueList> pageFetcher;
+      if (isArrow) {
+        long initialRowOffset = (long) firstPageRows.size();
+        Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
+        Number maxResultsOpt = (Number) optionsMap.get(BigQueryRpc.Option.MAX_RESULTS);
+        Long maxResults = maxResultsOpt != null ? maxResultsOpt.longValue() : null;
+        pageFetcher =
+            new ArrowQueryPageFetcher(
+                jobId, schema, arrowSchemaPojo, getOptions(), initialRowOffset, maxResults);
+      } else {
+        pageFetcher = new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options));
+      }
+
       return TableResult.newBuilder()
           .setSchema(schema)
           .setTotalRows(numRows)
-          .setPageNoSchema(
-              new PageImpl<>(
-                  // fetch next pages of results
-                  new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options)),
-                  cursor,
-                  transformTableData(
-                      results.getRows(),
-                      schema,
-                      getOptions().getDataFormatOptions().useInt64Timestamp())))
+          .setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows))
           .setJobId(jobId)
           .setQueryId(results.getQueryId())
           .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
-          .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+          .setRowsInPage((long) firstPageRows.size())
           .build();
     }
-    // only 1 page of result
+
     return TableResult.newBuilder()
         .setSchema(schema)
         .setTotalRows(numRows)
         .setPageNoSchema(
             new PageImpl<>(
-                new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
+                isArrow
+                    ? null
+                    : new TableDataPageFetcher(
+                        null, schema, getOptions(), null, optionMap(options)),
                 null,
-                transformTableData(
-                    results.getRows(),
-                    schema,
-                    getOptions().getDataFormatOptions().useInt64Timestamp())))
+                firstPageRows))
         // Return the JobID of the successful job
         .setJobId(
             results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
         .setQueryId(results.getQueryId())
         .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
-        .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+        .setRowsInPage((long) firstPageRows.size())
         .build();
   }
 
@@ -2207,6 +2457,10 @@ && getOptions().getOpenTelemetryTracer() != null) {
 
         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) {