From ab33e196b154478b94eee28312398fffdc71cf71 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Tue, 13 Oct 2026 03:52:27 +0000 Subject: [PATCH] perf(bigquery): [T1-12] replace Channels.newChannel bounce buffer with direct ByteBuffer channel on Arrow batches --- .../cloud/bigquery/ArrowDeserializer.java | 53 +++++++++++++++++-- .../cloud/bigquery/ArrowQueryResultImpl.java | 3 +- .../google/cloud/bigquery/ConnectionImpl.java | 7 +-- .../cloud/bigquery/ArrowDeserializerTest.java | 38 +++++++++++++ 4 files changed, 91 insertions(+), 10 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java index 08a2cb4612ff..f6cb627a34ef 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java +++ b/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) { diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java index 3f5d1b43e545..8808bc5dcf96 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java +++ b/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)); } /** diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java index b8832edb19de..ed7751fa3f65 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ConnectionImpl.java +++ b/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 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 vectors = new ArrayList<>(); List 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); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java index f39c8a0c0dfe..aed8ff3f80a2 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java +++ b/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))); + } }