perf(bigquery): [T1-12] replace Channels.newChannel bounce buffer with direct ByteBuffer channel on Arrow batches

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

@@ -21,9 +21,12 @@
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.io.BaseEncoding;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.math.BigDecimal;
-import java.nio.channels.Channels;
+import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
+import java.nio.channels.ReadableByteChannel;
 import java.nio.charset.StandardCharsets;
 import java.time.LocalDate;
 import java.time.LocalTime;
@@ -81,6 +84,51 @@ static BufferAllocator createChildAllocator(String name) {
 
   private ArrowDeserializer() {}
 
+  /**
+   * Returns a lightweight, unsynchronized {@link ReadableByteChannel} backed directly by the
+   * read-only {@link ByteBuffer} view of a protobuf {@link ByteString}, avoiding {@code
+   * Channels.newChannel(byteString.newInput())} transfer buffers and {@link
+   * ByteString#toByteArray()} array copies.
+   */
+  static ReadableByteChannel newByteStringChannel(ByteString byteString) {
+    return new ByteBufferReadableByteChannel(byteString.asReadOnlyByteBuffer());
+  }
+
+  private static final class ByteBufferReadableByteChannel implements ReadableByteChannel {
+    private final ByteBuffer buffer;
+    private boolean open = true;
+
+    ByteBufferReadableByteChannel(ByteBuffer buffer) {
+      this.buffer = buffer;
+    }
+
+    @Override
+    public int read(ByteBuffer dst) throws IOException {
+      if (!open) {
+        throw new ClosedChannelException();
+      }
+      if (!buffer.hasRemaining()) {
+        return -1;
+      }
+      int toRead = Math.min(dst.remaining(), buffer.remaining());
+      int oldLimit = buffer.limit();
+      buffer.limit(buffer.position() + toRead);
+      dst.put(buffer);
+      buffer.limit(oldLimit);
+      return toRead;
+    }
+
+    @Override
+    public boolean isOpen() {
+      return open;
+    }
+
+    @Override
+    public void close() {
+      open = false;
+    }
+  }
+
   /**
    * Deserializes a raw binary Arrow schema payload into an Apache Arrow Schema object.
    *
@@ -163,8 +211,7 @@ static boolean loadArrowRows(
           com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch =
               response.getArrowRecordBatch();
           try (ReadChannel readChannel =
-                  new ReadChannel(
-                      Channels.newChannel(batch.getSerializedRecordBatch().newInput()));
+                  new ReadChannel(newByteStringChannel(batch.getSerializedRecordBatch()));
               ArrowRecordBatch deserializedBatch =
                   MessageSerializer.deserializeRecordBatch(readChannel, childAllocator)) {
             if (deserializedBatch == null) {

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

@@ -23,7 +23,6 @@
 import com.google.cloud.bigquery.storage.v1.ReadSession;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
-import java.nio.channels.Channels;
 import java.nio.channels.ReadableByteChannel;
 import java.util.Collections;
 import java.util.Iterator;
@@ -526,7 +525,7 @@ private void loadBatch(byte[] bytes) throws IOException {
      * @throws IOException if deserialization fails
      */
     private void loadBatch(ByteString byteString) throws IOException {
-      loadBatch(Channels.newChannel(byteString.newInput()));
+      loadBatch(ArrowDeserializer.newByteStringChannel(byteString));
     }
 
     /**

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

@@ -80,7 +80,6 @@
 import org.apache.arrow.vector.ipc.ReadChannel;
 import org.apache.arrow.vector.ipc.message.MessageSerializer;
 import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
 
 /** Implementation for {@link Connection}, the generic BigQuery connection API (not JDBC). */
 class ConnectionImpl implements Connection {
@@ -1114,8 +1113,7 @@ private ArrowRowReader(ArrowSchema arrowSchema, Map<String, Integer> arrowNameTo
       org.apache.arrow.vector.types.pojo.Schema schema =
           MessageSerializer.deserializeSchema(
               new org.apache.arrow.vector.ipc.ReadChannel(
-                  new ByteArrayReadableSeekableByteChannel(
-                      arrowSchema.getSerializedSchema().toByteArray())));
+                  ArrowDeserializer.newByteStringChannel(arrowSchema.getSerializedSchema())));
       List<FieldVector> vectors = new ArrayList<>();
       List<Field> fields = schema.getFields();
       for (int i = 0; i < fields.size(); i++) {
@@ -1138,8 +1136,7 @@ private void processRows(
         org.apache.arrow.vector.ipc.message.ArrowRecordBatch deserializedBatch =
             MessageSerializer.deserializeRecordBatch(
                 new ReadChannel(
-                    new ByteArrayReadableSeekableByteChannel(
-                        batch.getSerializedRecordBatch().toByteArray())),
+                    ArrowDeserializer.newByteStringChannel(batch.getSerializedRecordBatch())),
                 allocator);
 
         loader.load(deserializedBatch);

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

@@ -899,4 +899,42 @@ public void testArrowVectorToFieldValue_structSubfieldMismatchThrows() {
       }
     }
   }
+
+  @Test
+  public void testNewByteStringChannelReadsIncrementalSlicesAndCloses() throws IOException {
+    byte[] source = "0123456789abcdef".getBytes(StandardCharsets.UTF_8);
+    ByteString byteString = ByteString.copyFrom(source);
+    java.nio.channels.ReadableByteChannel channel =
+        ArrowDeserializer.newByteStringChannel(byteString);
+    assertTrue(channel.isOpen());
+
+    java.nio.ByteBuffer dst = java.nio.ByteBuffer.allocate(6);
+    assertEquals(6, channel.read(dst));
+    dst.flip();
+    byte[] chunk1 = new byte[6];
+    dst.get(chunk1);
+    assertEquals("012345", new String(chunk1, StandardCharsets.UTF_8));
+
+    dst.clear();
+    assertEquals(6, channel.read(dst));
+    dst.flip();
+    byte[] chunk2 = new byte[6];
+    dst.get(chunk2);
+    assertEquals("6789ab", new String(chunk2, StandardCharsets.UTF_8));
+
+    dst.clear();
+    assertEquals(4, channel.read(dst));
+    dst.flip();
+    byte[] chunk3 = new byte[4];
+    dst.get(chunk3);
+    assertEquals("cdef", new String(chunk3, StandardCharsets.UTF_8));
+
+    dst.clear();
+    assertEquals(-1, channel.read(dst));
+
+    channel.close();
+    assertFalse(channel.isOpen());
+    assertThrows(
+        java.nio.channels.ClosedChannelException.class, () -> channel.read(java.nio.ByteBuffer.allocate(4)));
+  }
 }