refactor(bigquery): address PR 13943 review comments on channel lifecycle and schema resolution

Author: jinseopkim0

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

@@ -86,8 +86,11 @@ private ArrowDeserializer() {}
    * @throws IOException if deserialization of the Arrow schema fails
    */
   static Object deserializeSchema(byte[] schemaBytes) throws IOException {
-    return MessageSerializer.deserializeSchema(
-        new ReadChannel(new ByteArrayReadableSeekableByteChannel(schemaBytes)));
+    try (ByteArrayReadableSeekableByteChannel byteChannel =
+            new ByteArrayReadableSeekableByteChannel(schemaBytes);
+        ReadChannel readChannel = new ReadChannel(byteChannel)) {
+      return MessageSerializer.deserializeSchema(readChannel);
+    }
   }
 
   /**
@@ -121,12 +124,30 @@ static Object jsonToArrowSchema(String json) {
     }
   }
 
+  /**
+   * Resolves an Apache Arrow Schema from either an in-memory Schema POJO or a serialized JSON
+   * string.
+   *
+   * @param arrowSchema the Arrow schema POJO or JSON string representation
+   * @return the resolved Apache Arrow Schema, or null if schema cannot be resolved
+   * @throws IOException if parsing JSON fails
+   */
+  private static org.apache.arrow.vector.types.pojo.Schema resolveArrowSchema(Object arrowSchema)
+      throws IOException {
+    if (arrowSchema instanceof org.apache.arrow.vector.types.pojo.Schema) {
+      return (org.apache.arrow.vector.types.pojo.Schema) arrowSchema;
+    }
+    if (arrowSchema instanceof String) {
+      return org.apache.arrow.vector.types.pojo.Schema.fromJSON((String) arrowSchema);
+    }
+    return null;
+  }
+
   /**
    * Reads and decodes a batch of Arrow rows from the provided stream iterator into the row batch.
    *
    * @param iterator the stream iterator providing ReadRowsResponse messages
-   * @param arrowSchemaPojo the Arrow schema pojo (or null if restoring from json)
-   * @param arrowSchemaJson the Arrow schema JSON representation
+   * @param arrowSchema the Arrow schema POJO or serialized JSON representation
    * @param schema the BigQuery target Schema
    * @param rowBatch the destination list for decoded rows
    * @param pageSize the maximum number of rows to decode in this batch
@@ -137,28 +158,24 @@ static Object jsonToArrowSchema(String json) {
    */
   static boolean loadArrowRows(
       Iterator<ReadRowsResponse> iterator,
-      Object arrowSchemaPojo,
-      String arrowSchemaJson,
+      Object arrowSchema,
       Schema schema,
       List<FieldValueList> rowBatch,
       long pageSize,
       long totalRowsReturned,
       long maxResults)
       throws IOException {
-    org.apache.arrow.vector.types.pojo.Schema arrowSchema =
-        arrowSchemaPojo instanceof org.apache.arrow.vector.types.pojo.Schema
-            ? (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaPojo
-            : (arrowSchemaJson != null
-                ? org.apache.arrow.vector.types.pojo.Schema.fromJSON(arrowSchemaJson)
-                : null);
+    org.apache.arrow.vector.types.pojo.Schema resolvedSchema = resolveArrowSchema(arrowSchema);
 
-    if (arrowSchema == null) {
+    if (resolvedSchema == null) {
       return false;
     }
 
+    org.apache.arrow.vector.types.pojo.Schema arrowSchemaFinal = resolvedSchema;
+
     try (BufferAllocator childAllocator =
             AllocatorHolder.ALLOCATOR.newChildAllocator("loadArrowRows", 0, Long.MAX_VALUE);
-        VectorSchemaRoot closedRoot = createVectorSchemaRoot(arrowSchema, childAllocator)) {
+        VectorSchemaRoot closedRoot = createVectorSchemaRoot(arrowSchemaFinal, childAllocator)) {
       VectorLoader loader = new VectorLoader(closedRoot);
       boolean hasMore = false;
       while (rowBatch.size() < pageSize

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java

@@ -214,14 +214,7 @@ public void testLoadArrowRows_multiBatchStream() throws IOException {
       List<FieldValueList> rowBatch = new ArrayList<>();
       boolean hasMore =
           ArrowDeserializer.loadArrowRows(
-              Arrays.asList(r1, r2).iterator(),
-              arrowSchema,
-              null,
-              bqSchema,
-              rowBatch,
-              10L,
-              0L,
-              10L);
+              Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 10L, 0L, 10L);
 
       assertFalse(hasMore);
       assertEquals(4, rowBatch.size());
@@ -246,7 +239,7 @@ public void testLoadArrowRows_respectsPageSize() throws IOException {
       List<FieldValueList> rowBatch = new ArrayList<>();
       boolean hasMore =
           ArrowDeserializer.loadArrowRows(
-              Arrays.asList(r1, r2).iterator(), arrowSchema, null, bqSchema, rowBatch, 2L, 0L, 10L);
+              Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 2L, 0L, 10L);
 
       assertTrue(hasMore);
       assertEquals(2, rowBatch.size());
@@ -269,7 +262,7 @@ public void testLoadArrowRows_respectsMaxResults() throws IOException {
       List<FieldValueList> rowBatch = new ArrayList<>();
       boolean hasMore =
           ArrowDeserializer.loadArrowRows(
-              Arrays.asList(r1, r2).iterator(), arrowSchema, null, bqSchema, rowBatch, 10L, 0L, 3L);
+              Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 10L, 0L, 3L);
 
       assertFalse(hasMore);
       assertEquals(3, rowBatch.size());
@@ -294,7 +287,7 @@ public void testLoadArrowRows_unconsumedBatchRowsSignalHasMore() throws IOExcept
       List<FieldValueList> rowBatch = new ArrayList<>();
       boolean hasMore =
           ArrowDeserializer.loadArrowRows(
-              Arrays.asList(r1).iterator(), arrowSchema, null, bqSchema, rowBatch, 2L, 0L, 10L);
+              Arrays.asList(r1).iterator(), arrowSchema, bqSchema, rowBatch, 2L, 0L, 10L);
 
       assertTrue(hasMore);
       assertEquals(2, rowBatch.size());
@@ -310,7 +303,6 @@ public void testLoadArrowRows_nullSchemaReturnsFalse() throws IOException {
         ArrowDeserializer.loadArrowRows(
             Arrays.<ReadRowsResponse>asList().iterator(),
             null,
-            null,
             Schema.of(),
             rowBatch,
             10L,