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

[jinseopkim0]

Stacked PR 3 of 3: Integrates the Arrow API surface (PR 1) and the ArrowDeserializer (PR 2) into the Veneer client's query execution path, implementing fallback validations, first-page parsing, and stateful gRPC stream pagination.

[gemini-code-assist]

Code Review

This pull request introduces support for Arrow-formatted query results in the BigQuery client, adding ArrowQueryPageFetcher to fetch Arrow pages and updating response parsing to handle Arrow schemas and record batches. The review feedback highlights critical issues that need to be addressed: a potential NotSerializableException in ArrowQueryPageFetcher due to a non-serializable schema field, memory leaks from unclosed Arrow vectors and record batches, and overly restrictive type checks (instanceof List instead of instanceof Collection) when determining page row counts.

[gemini-code-assist]

NextPageFetcher is a Serializable interface, meaning implementations like ArrowQueryPageFetcher must be fully serializable. However, org.apache.arrow.vector.types.pojo.Schema does not implement Serializable, which will cause a NotSerializableException if the page or fetcher is serialized.

To fix this, we can store the schema as a JSON string (arrowSchemaJson) using Arrow's built-in Schema.toJson() and Schema.fromJSON(String) methods, and mark the arrowSchemaPojo field as transient so it is lazily deserialized when needed.

    private final JobId jobId;
    private final Schema schema;
    private final String arrowSchemaJson;
    private final BigQueryOptions serviceOptions;
    private final long maxResults;

    private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo;
    private transient BigQueryReadClient bqReadClient;
    private transient ServerStream<ReadRowsResponse> stream;
    private transient Iterator<ReadRowsResponse> streamIterator;
    private long totalRowsReturned = 0L;
    private boolean streamClosed = false;

    ArrowQueryPageFetcher(
        JobId jobId,
        Schema schema,
        org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo,
        BigQueryOptions serviceOptions,
        long initialRowOffset,
        Long maxResults) {
      this.jobId = jobId;
      this.schema = schema;
      this.arrowSchemaJson = arrowSchemaPojo != null ? arrowSchemaPojo.toJson() : null;
      this.arrowSchemaPojo = arrowSchemaPojo;
      this.serviceOptions = serviceOptions;
      this.totalRowsReturned = initialRowOffset;
      this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE;
    }

References

1. Prefer lazy initialization over eager initialization for resource-intensive objects (such as CharsetEncoder) if they are not guaranteed to be used in all execution paths, to avoid unnecessary performance and memory overhead.

[gemini-code-assist]

There are two critical resource management issues in this block:

1. Vector Allocation Leak: If an exception occurs while creating vectors in the loop, any previously allocated FieldVector instances in the vectors list will leak because they are not closed. Using VectorSchemaRoot.create(arrowSchemaPojo, allocator) is the standard, safe way to instantiate a VectorSchemaRoot and automatically handles cleanup of all vectors if allocation fails.
2. ArrowRecordBatch Leak: org.apache.arrow.vector.ipc.message.ArrowRecordBatch is a closeable resource. If loader.load(deserializedBatch) throws an exception, deserializedBatch.close() is bypassed, causing a memory leak. Wrapping it in a try-with-resources block ensures it is always closed safely.

Additionally, we integrate the lazy deserialization of arrowSchemaPojo from arrowSchemaJson to support serialization.

        if (arrowSchemaPojo == null && arrowSchemaJson != null) {
          arrowSchemaPojo = org.apache.arrow.vector.types.pojo.Schema.fromJSON(arrowSchemaJson);
        }

        try (org.apache.arrow.memory.BufferAllocator allocator =
            new org.apache.arrow.memory.RootAllocator(Long.MAX_VALUE);
            org.apache.arrow.vector.VectorSchemaRoot root =
                org.apache.arrow.vector.VectorSchemaRoot.create(arrowSchemaPojo, allocator)) {
          org.apache.arrow.vector.VectorLoader loader =
              new org.apache.arrow.vector.VectorLoader(root);

          while (rowBatch.size() < pageSize && streamIterator.hasNext()) {
            ReadRowsResponse response = streamIterator.next();
            if (response.hasArrowRecordBatch()) {
              com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch =
                  response.getArrowRecordBatch();
              try (org.apache.arrow.vector.ipc.message.ArrowRecordBatch deserializedBatch =
                  org.apache.arrow.vector.ipc.message.MessageSerializer.deserializeRecordBatch(
                      new org.apache.arrow.vector.ipc.ReadChannel(
                          new org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel(
                              batch.getSerializedRecordBatch().toByteArray())),
                      allocator)) {
                loader.load(deserializedBatch);
                int batchRowCount = root.getRowCount();
                for (int i = 0; i < batchRowCount; i++) {
                  rowBatch.add(ArrowDeserializer.arrowRootToFieldValueList(root, i, schema));
                }
                root.clear();
              }
            }
          }
        }

References

1. Prefer lazy initialization over eager initialization for resource-intensive objects (such as CharsetEncoder) if they are not guaranteed to be used in all execution paths, to avoid unnecessary performance and memory overhead.

[gemini-code-assist]

Checking if firstPageRows is an instance of List is restrictive. If ArrowDeserializer.deserializeRecordBatch returns a different type of collection (e.g., a custom Collection or Set), the check will fail and initialRowOffset will default to 0L. Checking for java.util.Collection is safer and more robust.

Suggested change

	long initialRowOffset =
	firstPageRows instanceof List ? ((List<?>) firstPageRows).size() : 0L;
	long initialRowOffset =
	firstPageRows instanceof java.util.Collection ? ((java.util.Collection<?>) firstPageRows).size() : 0L;

[gemini-code-assist]

Checking if firstPageRows is an instance of List is restrictive. If ArrowDeserializer.deserializeRecordBatch returns a different type of collection (e.g., a custom Collection or Set), the check will fail and rowsInPage will default to 0L. Checking for java.util.Collection is safer and more robust.

Suggested change

	.setRowsInPage(
	firstPageRows instanceof List ? (long) ((List<?>) firstPageRows).size() : 0L)
	.setRowsInPage(
	firstPageRows instanceof java.util.Collection ? (long) ((java.util.Collection<?>) firstPageRows).size() : 0L)

[gemini-code-assist]

Checking if firstPageRows is an instance of List is restrictive. If ArrowDeserializer.deserializeRecordBatch returns a different type of collection (e.g., a custom Collection or Set), the check will fail and rowsInPage will default to 0L. Checking for java.util.Collection is safer and more robust.

Suggested change

	.setRowsInPage(firstPageRows instanceof List ? (long) ((List<?>) firstPageRows).size() : 0L)
	.setRowsInPage(firstPageRows instanceof java.util.Collection ? (long) ((java.util.Collection<?>) firstPageRows).size() : 0L)

[jinseopkim0]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request introduces support for the Apache Arrow query results format in the BigQuery client. It refactors Arrow schema conversion and vector creation logic into a new helper class ArrowPojoUtils, adds an ArrowQueryPageFetcher to stream Arrow-formatted query results, and updates BigQueryImpl and QueryRequestInfo to handle Arrow serialization options and deserialize Arrow record batches. The review feedback highlights several improvement opportunities: preventing a potential resource leak in ArrowPojoUtils.createVectors by closing already allocated vectors if an exception is thrown, avoiding a potential NullPointerException in BigQueryImpl by using serviceOptions.getCredentialsProvider(), and extracting duplicate logic for computing firstPageRows to improve maintainability.

[gemini-code-assist]

If an exception is thrown during field.createVector(allocator) (for example, due to allocator capacity limits or other runtime issues), any vectors already added to the vectors list will not be closed, leading to a memory/resource leak. Wrapping the loop in a try-catch block to close any already allocated vectors in the catch block before rethrowing the exception prevents this leak.

  static List<FieldVector> createVectors(Object arrowSchemaObj, BufferAllocator allocator) {
    org.apache.arrow.vector.types.pojo.Schema arrowSchema =
        (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj;
    List<FieldVector> vectors = new ArrayList<>();
    try {
      for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) {
        vectors.add(field.createVector(allocator));
      }
      return vectors;
    } catch (Throwable t) {
      for (int i = vectors.size() - 1; i >= 0; i--) {
        try {
          vectors.get(i).close();
        } catch (Exception e) {
          // Suppress exception during cleanup to preserve original exception
        }
      }
      throw t;
    }
  }

References

1. When managing a collection of closeable resources (e.g., scopes), ensure they are closed in the reverse order of their creation (LIFO). The implementation must be exception-safe to prevent resource leaks, meaning all opened resources should be closed even if exceptions occur during their creation or closing.

[gemini-code-assist]

If serviceOptions.getCredentials() is null (e.g., when running with local emulators or custom credentials setups), FixedCredentialsProvider.create(null) will throw a NullPointerException. Using serviceOptions.getCredentialsProvider() is safer and more robust, as it directly retrieves the configured credentials provider and avoids this issue. This also allows you to remove the unused FixedCredentialsProvider import.

        if (bqReadClient == null) {
          BigQueryReadSettings settings =
              BigQueryReadSettings.newBuilder()
                  .setCredentialsProvider(serviceOptions.getCredentialsProvider())
                  .build();
          bqReadClient = BigQueryReadClient.create(settings);
        }

[gemini-code-assist]

The logic to compute firstPageRows is identical in both branches of the if (results.getPageToken() != null) check. Extracting this logic before the check eliminates a significant amount of duplicate code, improving readability and maintainability.

    Iterable<FieldValueList> firstPageRows;
    if (isArrow) {
      if (results.getArrowRecordBatch() != null) {
        try {
          firstPageRows =
              ArrowDeserializer.deserializeRecordBatch(
                  results.getArrowRecordBatch().decodeSerializedRecordBatch(),
                  schema,
                  arrowSchemaPojo);
        } catch (IOException e) {
          throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e);
        }
      } else {
        firstPageRows = ImmutableList.of();
      }
    } else {
      firstPageRows =
          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 =
            firstPageRows instanceof Collection ? ((Collection<?>) firstPageRows).size() : 0L;
        pageFetcher =
            new ArrowQueryPageFetcher(
                jobId,
                schema,
                (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaPojo,
                getOptions(),
                initialRowOffset,
                null); // Or use maxResults from configuration if available
      } else {
        pageFetcher = new QueryPageFetcher(jobId, schema, getOptions(), cursor, 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(
              firstPageRows instanceof Collection
                  ? (long) ((Collection<?>) firstPageRows).size()
                  : 0L)
          .build();
    }

    return TableResult.newBuilder()
        .setSchema(schema)
        .setTotalRows(numRows)
        .setPageNoSchema(
            new PageImpl<>(
                isArrow
                    ? null
                    : new TableDataPageFetcher(
                        null, schema, getOptions(), null, optionMap(options)),
                null,
                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(
            firstPageRows instanceof Collection
                ? (long) ((Collection<?>) firstPageRows).size()
                : 0L)
        .build();

[jinseopkim0]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request introduces support for the Apache Arrow query results format in the BigQuery client. It extracts Arrow conversion utilities into a new ArrowPojoUtils helper class, implements ArrowQueryPageFetcher to stream results via the BigQuery Storage Read API, and updates query execution paths to handle Arrow schemas and record batches. The review feedback focuses on critical resource management and performance optimizations, specifically ensuring exception-safe LIFO cleanup of Arrow vectors, allocators, and schema roots to prevent memory leaks, reusing allocator and root instances across page fetches, enforcing the maxResults limit, and properly handling credentials and universe domains.

[gemini-code-assist]

Wrap the vector creation loop in a try-catch block to close any successfully created vectors in reverse order (LIFO) if an exception is thrown during the process. This ensures exception safety and prevents resource leaks.

  static List<FieldVector> createVectors(Object arrowSchemaObj, BufferAllocator allocator) {
    org.apache.arrow.vector.types.pojo.Schema arrowSchema =
        (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj;
    List<FieldVector> vectors = new ArrayList<>();
    try {
      for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) {
        vectors.add(field.createVector(allocator));
      }
      return vectors;
    } catch (Throwable t) {
      for (int i = vectors.size() - 1; i >= 0; i--) {
        try {
          vectors.get(i).close();
        } catch (Exception e) {
          // ignore
        }
      }
      throw t;
    }
  }

References

1. When managing a collection of closeable resources, ensure they are closed in the reverse order of their creation (LIFO). The implementation must be exception-safe to prevent resource leaks, meaning all opened resources should be closed even if exceptions occur during their creation or closing.

[gemini-code-assist]

Reuse the transient allocator and root fields instead of creating them on every page fetch. Additionally, enforce the maxResults limit inside the record batch loading loop to prevent returning more rows than requested.

        if (allocator == null) {
          allocator = new RootAllocator(Long.MAX_VALUE);
        }
        if (root == null) {
          root = VectorSchemaRoot.create(arrowSchemaPojo, allocator);
        }

        VectorLoader loader = new VectorLoader(root);

        while (rowBatch.size() < pageSize && streamIterator.hasNext() && (totalRowsReturned + rowBatch.size() < maxResults)) {
          ReadRowsResponse response = streamIterator.next();
          if (response.hasArrowRecordBatch()) {
            com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch =
                response.getArrowRecordBatch();
            try (ArrowRecordBatch deserializedBatch =
                MessageSerializer.deserializeRecordBatch(
                    new ReadChannel(
                        new ByteArrayReadableSeekableByteChannel(
                            batch.getSerializedRecordBatch().toByteArray())),
                    allocator)) {
              loader.load(deserializedBatch);
              int batchRowCount = root.getRowCount();
              for (int i = 0; i < batchRowCount; i++) {
                if (totalRowsReturned + rowBatch.size() >= maxResults) {
                  break;
                }
                rowBatch.add(ArrowDeserializer.arrowRootToFieldValueList(root, i, schema));
              }
              root.clear();
            }
          }
        }

[gemini-code-assist]

Declare allocator and root as transient fields in ArrowQueryPageFetcher so they can be reused across multiple calls to getNextPage(), avoiding the overhead of creating a new RootAllocator and VectorSchemaRoot on every page fetch.

    private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo;
    private transient BigQueryReadClient bqReadClient;
    private transient ServerStream<ReadRowsResponse> stream;
    private transient Iterator<ReadRowsResponse> streamIterator;
    private transient BufferAllocator allocator;
    private transient VectorSchemaRoot root;
    private long totalRowsReturned = 0L;
    private boolean streamClosed = false;

[gemini-code-assist]

Ensure that serviceOptions.getCredentials() is not null before setting the credentials provider, and propagate the custom universe domain if configured in serviceOptions to support VPC-SC and custom endpoints.

        if (bqReadClient == null) {
          BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder();
          if (serviceOptions.getCredentials() != null) {
            settingsBuilder.setCredentialsProvider(
                FixedCredentialsProvider.create(serviceOptions.getCredentials()));
          }
          if (serviceOptions.getUniverseDomain() != null) {
            settingsBuilder.setUniverseDomain(serviceOptions.getUniverseDomain());
          }
          bqReadClient = BigQueryReadClient.create(settingsBuilder.build());
        }

[gemini-code-assist]

Close the transient root, allocator, and bqReadClient resources in the reverse order of their creation (LIFO) and in an exception-safe manner to prevent resource leaks.

    private void closeClient() {
      if (root != null) {
        try {
          root.close();
        } catch (Exception e) {
          // ignore
        } finally {
          root = null;
        }
      }
      if (allocator != null) {
        try {
          allocator.close();
        } catch (Exception e) {
          // ignore
        } finally {
          allocator = null;
        }
      }
      if (bqReadClient != null) {
        try {
          bqReadClient.close();
        } catch (Exception e) {
          // ignore
        } finally {
          bqReadClient = null;
        }
      }
      streamIterator = null;
      stream = null;
    }

References

1. When managing a collection of closeable resources, ensure they are closed in the reverse order of their creation (LIFO). The implementation must be exception-safe to prevent resource leaks, meaning all opened resources should be closed even if exceptions occur during their creation or closing.