fix(bigquery): address gemini-code-assist review feedback on Arrow page fetcher and exception-safe vector creation

Author: jinseopkim0

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

@@ -83,10 +83,21 @@ static List<FieldVector> createVectors(Object arrowSchemaObj, BufferAllocator al
     org.apache.arrow.vector.types.pojo.Schema arrowSchema =
         (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj;
     List<FieldVector> vectors = new ArrayList<>();
-    for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) {
-      vectors.add(field.createVector(allocator));
+    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;
     }
-    return vectors;
   }
 
   private static LegacySQLTypeName arrowTypeToLegacySQLTypeName(ArrowType type) {

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

@@ -293,6 +293,8 @@ private static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValue
     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;
 
@@ -324,12 +326,15 @@ public Page<FieldValueList> getNextPage() {
 
       try {
         if (bqReadClient == null) {
-          BigQueryReadSettings settings =
-              BigQueryReadSettings.newBuilder()
-                  .setCredentialsProvider(
-                      FixedCredentialsProvider.create(serviceOptions.getCredentials()))
-                  .build();
-          bqReadClient = BigQueryReadClient.create(settings);
+          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());
         }
 
         if (streamIterator == null) {
@@ -354,28 +359,37 @@ public Page<FieldValueList> getNextPage() {
           arrowSchemaPojo = org.apache.arrow.vector.types.pojo.Schema.fromJSON(arrowSchemaJson);
         }
 
-        try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
-            VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchemaPojo, allocator)) {
-          VectorLoader loader = new 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 (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++) {
-                  rowBatch.add(ArrowDeserializer.arrowRootToFieldValueList(root, i, schema));
+        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;
                 }
-                root.clear();
+                rowBatch.add(ArrowDeserializer.arrowRootToFieldValueList(root, i, schema));
               }
+              root.clear();
             }
           }
         }
@@ -406,9 +420,32 @@ public Page<FieldValueList> getNextPage() {
     }
 
     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) {
-        bqReadClient.close();
-        bqReadClient = null;
+        try {
+          bqReadClient.close();
+        } catch (Exception e) {
+          // ignore
+        } finally {
+          bqReadClient = null;
+        }
       }
       streamIterator = null;
       stream = null;