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)