From 8e6ccae0efbba03f31b9639a1a387da3255d932a Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Sat, 18 Jul 2026 16:39:39 +0200 Subject: [PATCH 1/8] Initial commit --- .../arrow/flight/grpc/GetReadableBuffer.java | 122 ++++++----- .../flight/grpc/TestGetReadableBuffer.java | 201 ++++++++++++++++++ 2 files changed, 267 insertions(+), 56 deletions(-) create mode 100644 flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index fcba88d212..dd519b7b40 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -16,84 +16,94 @@ */ package org.apache.arrow.flight.grpc; -import com.google.common.base.Throwables; import com.google.common.io.ByteStreams; -import io.grpc.internal.ReadableBuffer; +import io.grpc.HasByteBuffer; import java.io.IOException; import java.io.InputStream; -import java.lang.reflect.Field; +import java.nio.ByteBuffer; import org.apache.arrow.memory.ArrowBuf; /** - * Enable access to ReadableBuffer directly to copy data from a BufferInputStream into a target - * ByteBuffer/ByteBuf. + * Copy data from a gRPC-provided InputStream into a target ArrowBuf. * - *

This could be solved by BufferInputStream exposing Drainable. + *

When the stream is backed by gRPC's own buffers (i.e. it implements {@link HasByteBuffer}), + * the payload is copied directly out of gRPC's {@link ByteBuffer}s, avoiding an intermediate heap + * {@code byte[]} allocation and the extra copy that goes with it. Otherwise, we fall back to + * reading the stream into a heap array. + * + *

This relies on gRPC's public zero-copy APIs ({@link HasByteBuffer#getByteBuffer()} and + * {@link InputStream#skip(long)}). No longer access gRPC internals via reflection. See + * apache/arrow-java#939. */ -public class GetReadableBuffer { - - private static final Field READABLE_BUFFER; - private static final Class BUFFER_INPUT_STREAM; - - static { - Field tmpField = null; - Class tmpClazz = null; - try { - Class clazz = Class.forName("io.grpc.internal.ReadableBuffers$BufferInputStream"); +public final class GetReadableBuffer { - Field f = clazz.getDeclaredField("buffer"); - f.setAccessible(true); - // don't set until we've gotten past all exception cases. - tmpField = f; - tmpClazz = clazz; - } catch (Exception e) { - new RuntimeException("Failed to initialize GetReadableBuffer, falling back to slow path", e) - .printStackTrace(); - } - READABLE_BUFFER = tmpField; - BUFFER_INPUT_STREAM = tmpClazz; - } - - /** - * Extracts the ReadableBuffer for the given input stream. - * - * @param is Must be an instance of io.grpc.internal.ReadableBuffers$BufferInputStream or null - * will be returned. - */ - public static ReadableBuffer getReadableBuffer(InputStream is) { - - if (BUFFER_INPUT_STREAM == null || !is.getClass().equals(BUFFER_INPUT_STREAM)) { - return null; - } - - try { - return (ReadableBuffer) READABLE_BUFFER.get(is); - } catch (Exception ex) { - throw Throwables.propagate(ex); - } - } + private GetReadableBuffer() {} /** * Helper method to read a gRPC-provided InputStream into an ArrowBuf. * - * @param stream The stream to read from. Should be an instance of {@link #BUFFER_INPUT_STREAM}. + * @param stream The stream to read from. * @param buf The buffer to read into. * @param size The number of bytes to read. - * @param fastPath Whether to enable the fast path (i.e. detect whether the stream is a {@link - * #BUFFER_INPUT_STREAM}). - * @throws IOException if there is an error reading form the stream + * @param fastPath Whether to enable the fast path (i.e. copy directly from the stream's backing + * {@link ByteBuffer}s when the stream supports it). + * @throws IOException if there is an error reading from the stream */ public static void readIntoBuffer( final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath) throws IOException { - ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null; - byte[] heapBytes = new byte[size]; - if (readableBuffer != null) { - readableBuffer.readBytes(heapBytes, 0, size); + if (fastPath + && stream instanceof HasByteBuffer + && ((HasByteBuffer) stream).byteBufferSupported()) { + readFromByteBuffers((HasByteBuffer) stream, stream, buf, size); } else { + final byte[] heapBytes = new byte[size]; ByteStreams.readFully(stream, heapBytes); + buf.writeBytes(heapBytes); } - buf.writeBytes(heapBytes); buf.writerIndex(size); } + + /** + * Copy {@code size} bytes directly out of the stream's backing {@link ByteBuffer}s into {@code + * buf}. + * + *

{@link HasByteBuffer#getByteBuffer()} exposes the next chunk of readable bytes without + * advancing the stream, so we copy the chunk and then {@link InputStream#skip(long)} past (and + * release) the bytes we consumed before asking for the next chunk. + */ + private static void readFromByteBuffers( + final HasByteBuffer hasByteBuffer, + final InputStream stream, + final ArrowBuf buf, + final int size) + throws IOException { + long writeIndex = 0; + int remaining = size; + while (remaining > 0) { + final ByteBuffer chunk = hasByteBuffer.getByteBuffer(); + // getByteBuffer() may expose more than we need (e.g. bytes belonging to the following + // field), so only copy up to the number of bytes still outstanding. We are allowed to change + // the returned buffer's position/limit without affecting the stream. + final int toRead = chunk == null ? 0 : Math.min(remaining, chunk.remaining()); + if (toRead == 0) { + throw new IOException( + String.format( + "Unexpected end of stream: %d of %d bytes remaining to be read", remaining, size)); + } + chunk.limit(chunk.position() + toRead); + buf.setBytes(writeIndex, chunk); + // getByteBuffer() does not advance the stream; skip() consumes (and frees) the bytes. + long skipped = 0; + while (skipped < toRead) { + final long n = stream.skip(toRead - skipped); + if (n <= 0) { + throw new IOException("Failed to skip past consumed bytes in the gRPC stream"); + } + skipped += n; + } + writeIndex += toRead; + remaining -= toRead; + } + } } diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java new file mode 100644 index 0000000000..b6c92839bf --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight.grpc; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import io.grpc.HasByteBuffer; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.util.ArrayDeque; +import java.util.Arrays; +import java.util.Deque; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for reading a gRPC-provided {@link InputStream} into an {@link ArrowBuf}. */ +public class TestGetReadableBuffer { + + private BufferAllocator allocator; + + @BeforeEach + public void setUp() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void tearDown() { + allocator.close(); + } + + @Test + public void testFastPathSingleChunk() throws IOException { + final byte[] payload = payload(64); + try (ChunkedStream stream = new ChunkedStream(payload); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); + assertEquals(payload.length, buf.writerIndex()); + assertArrayEquals(payload, toBytes(buf, payload.length)); + } + } + + /** Copy loop stitches several backing buffers of the requested size. */ + @Test + public void testFastPathAcrossChunks() throws IOException { + final byte[] payload = payload(70); + try (ChunkedStream stream = new ChunkedStream(slice(payload, 10, 1, 32, 27)); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); + assertArrayEquals(payload, toBytes(buf, payload.length)); + } + } + + /** + * A chunk may hold more than the caller asked for (like next field bytes). Only {@code + * size} bytes may be consumed, the rest must remain readable. + */ + @Test + public void testFastPathDoesNotOverConsume() throws IOException { + final byte[] payload = payload(48); + try (ChunkedStream stream = new ChunkedStream(payload); + ArrowBuf buf = allocator.buffer(20)) { + GetReadableBuffer.readIntoBuffer(stream, buf, 20, true); + assertArrayEquals(Arrays.copyOf(payload, 20), toBytes(buf, 20)); + + final byte[] rest = new byte[payload.length - 20]; + assertEquals(rest.length, stream.read(rest)); + assertArrayEquals(Arrays.copyOfRange(payload, 20, payload.length), rest); + } + } + + /** A truncated stream must fail loudly rather than leave the buffer partially filled. */ + @Test + public void testFastPathTruncatedStream() throws IOException { + try (ChunkedStream stream = new ChunkedStream(payload(10)); + ArrowBuf buf = allocator.buffer(32)) { + assertThrows(IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 32, true)); + } + } + + /** Both take the heap-array path: streams without zero-copy support, and fastPath=false. */ + @Test + public void testSlowPath() throws IOException { + final byte[] payload = payload(33); + try (ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer( + new ByteArrayInputStream(payload), buf, payload.length, true); + assertArrayEquals(payload, toBytes(buf, payload.length)); + } + try (ChunkedStream stream = new ChunkedStream(payload); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, false); + assertArrayEquals(payload, toBytes(buf, payload.length)); + } + } + + private static byte[] payload(int size) { + final byte[] bytes = new byte[size]; + for (int i = 0; i < size; i++) { + bytes[i] = (byte) i; + } + return bytes; + } + + private static byte[] toBytes(ArrowBuf buf, int size) { + final byte[] bytes = new byte[size]; + buf.getBytes(0, bytes); + return bytes; + } + + private static byte[][] slice(byte[] payload, int... sizes) { + final byte[][] chunks = new byte[sizes.length][]; + int offset = 0; + for (int i = 0; i < sizes.length; i++) { + chunks[i] = Arrays.copyOfRange(payload, offset, offset + sizes[i]); + offset += sizes[i]; + } + return chunks; + } + + /** + * For gRPC's buffer-backed streams: {@link #getByteBuffer()} exposes the next chunk + * and {@link #skip(long)} advances. + */ + private static final class ChunkedStream extends InputStream implements HasByteBuffer { + private final Deque chunks = new ArrayDeque<>(); + + ChunkedStream(byte[]... chunks) { + for (byte[] chunk : chunks) { + this.chunks.add(ByteBuffer.wrap(chunk)); + } + } + + @Override + public boolean byteBufferSupported() { + return true; + } + + @Override + public ByteBuffer getByteBuffer() { + final ByteBuffer head = chunks.peek(); + // Like gRPC, hand out an independent view so the caller may adjust position/limit freely. + return head == null ? null : head.duplicate(); + } + + @Override + public long skip(long n) { + final ByteBuffer head = chunks.peek(); + if (head == null) { + return 0; + } + final int skipped = (int) Math.min(n, head.remaining()); + head.position(head.position() + skipped); + if (!head.hasRemaining()) { + chunks.poll(); + } + return skipped; + } + + @Override + public int read() { + final byte[] one = new byte[1]; + return read(one, 0, 1) == 1 ? one[0] & 0xFF : -1; + } + + @Override + public int read(byte[] dst, int off, int len) { + final ByteBuffer head = chunks.peek(); + if (head == null) { + return -1; + } + final int read = Math.min(len, head.remaining()); + head.get(dst, off, read); + if (!head.hasRemaining()) { + chunks.poll(); + } + return read; + } + } +} From 21b292d0c2e6d472b4407d6dafa5843593152da9 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Wed, 22 Jul 2026 23:39:38 +0200 Subject: [PATCH 2/8] Harden skip and regression test --- .../arrow/flight/grpc/GetReadableBuffer.java | 11 +++++++--- .../flight/grpc/TestGetReadableBuffer.java | 21 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index dd519b7b40..c925fbddf8 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -97,10 +97,15 @@ private static void readFromByteBuffers( long skipped = 0; while (skipped < toRead) { final long n = stream.skip(toRead - skipped); - if (n <= 0) { - throw new IOException("Failed to skip past consumed bytes in the gRPC stream"); + if (n > 0) { + skipped += n; + } else if (stream.read() == -1) { + throw new IOException("Unexpected end of stream while consuming copied bytes"); + } else { + // InputStream.skip() is permitted to return zero without reaching EOF. We have already + // copied this byte, so read it only to guarantee that the stream makes progress. + skipped++; } - skipped += n; } writeIndex += toRead; remaining -= toRead; diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java index b6c92839bf..dd21ca699d 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java @@ -90,6 +90,17 @@ public void testFastPathDoesNotOverConsume() throws IOException { } } + /** The copy loop must progress when a valid InputStream returns zero from skip(). */ + @Test + public void testFastPathSkipReturnsZero() throws IOException { + final byte[] payload = payload(32); + try (ChunkedStream stream = new ChunkedStream(true, payload); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); + assertArrayEquals(payload, toBytes(buf, payload.length)); + } + } + /** A truncated stream must fail loudly rather than leave the buffer partially filled. */ @Test public void testFastPathTruncatedStream() throws IOException { @@ -145,8 +156,14 @@ private static byte[][] slice(byte[] payload, int... sizes) { */ private static final class ChunkedStream extends InputStream implements HasByteBuffer { private final Deque chunks = new ArrayDeque<>(); + private boolean skipReturnsZero; ChunkedStream(byte[]... chunks) { + this(false, chunks); + } + + ChunkedStream(boolean skipReturnsZero, byte[]... chunks) { + this.skipReturnsZero = skipReturnsZero; for (byte[] chunk : chunks) { this.chunks.add(ByteBuffer.wrap(chunk)); } @@ -166,6 +183,10 @@ public ByteBuffer getByteBuffer() { @Override public long skip(long n) { + if (skipReturnsZero) { + skipReturnsZero = false; + return 0; + } final ByteBuffer head = chunks.peek(); if (head == null) { return 0; From 689790598b3484423f8cc8ca607a028070bb8d8d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sat, 25 Jul 2026 22:01:58 +0200 Subject: [PATCH 3/8] add detach path --- flight/flight-core/pom.xml | 1 + .../src/main/java/module-info.java | 1 - .../org/apache/arrow/flight/ArrowMessage.java | 180 ++++++++++++++- .../org/apache/arrow/flight/FlightStream.java | 1 + .../arrow/flight/grpc/GetReadableBuffer.java | 9 +- .../flight/TestArrowMessageDetachable.java | 211 ++++++++++++++++++ .../flight/grpc/TestGetReadableBuffer.java | 13 +- 7 files changed, 393 insertions(+), 23 deletions(-) create mode 100644 flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml index 905db61af7..ac05fa11f0 100644 --- a/flight/flight-core/pom.xml +++ b/flight/flight-core/pom.xml @@ -61,6 +61,7 @@ under the License. io.grpc grpc-core + runtime io.grpc diff --git a/flight/flight-core/src/main/java/module-info.java b/flight/flight-core/src/main/java/module-info.java index 669797ac93..9bafc5fddf 100644 --- a/flight/flight-core/src/main/java/module-info.java +++ b/flight/flight-core/src/main/java/module-info.java @@ -30,7 +30,6 @@ requires com.google.protobuf; requires com.google.protobuf.util; requires io.grpc; - requires io.grpc.internal; requires io.grpc.netty; requires io.grpc.protobuf; requires io.grpc.stub; diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java index ab4eab3048..816a380b04 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java @@ -23,7 +23,10 @@ import com.google.protobuf.CodedInputStream; import com.google.protobuf.CodedOutputStream; import com.google.protobuf.WireFormat; +import io.grpc.Detachable; import io.grpc.Drainable; +import io.grpc.HasByteBuffer; +import io.grpc.KnownLength; import io.grpc.MethodDescriptor.Marshaller; import io.grpc.protobuf.ProtoUtils; import io.netty.buffer.ByteBuf; @@ -46,6 +49,8 @@ import org.apache.arrow.flight.impl.Flight.FlightDescriptor; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.ForeignAllocation; +import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; @@ -280,12 +285,22 @@ public Iterable getBufs() { } private static ArrowMessage frame(BufferAllocator allocator, final InputStream stream) { + if (ENABLE_ZERO_COPY_READ) { + final ArrowBufInputStream detached = ArrowBufInputStream.tryCreate(allocator, stream); + if (detached != null) { + try { + return frame(allocator, detached); + } finally { + detached.close(); + } + } + } + FlightDescriptor descriptor = null; + MessageMetadataResult header = null; + ArrowBuf body = null; + ArrowBuf appMetadata = null; try { - FlightDescriptor descriptor = null; - MessageMetadataResult header = null; - ArrowBuf body = null; - ArrowBuf appMetadata = null; while (stream.available() > 0) { final int tagFirstByte = stream.read(); if (tagFirstByte == -1) { @@ -311,20 +326,22 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s } case APP_METADATA_TAG: { + if (appMetadata != null) { + appMetadata.close(); + appMetadata = null; + } int size = readRawVarint32(stream); - appMetadata = allocator.buffer(size); - GetReadableBuffer.readIntoBuffer(stream, appMetadata, size, ENABLE_ZERO_COPY_READ); + appMetadata = readBuffer(allocator, stream, size); break; } case BODY_TAG: if (body != null) { // only read last body. - body.getReferenceManager().release(); + body.close(); body = null; } int size = readRawVarint32(stream); - body = allocator.buffer(size); - GetReadableBuffer.readIntoBuffer(stream, body, size, ENABLE_ZERO_COPY_READ); + body = readBuffer(allocator, stream, size); break; default: @@ -362,9 +379,18 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s break; } } - return new ArrowMessage(descriptor, header, appMetadata, body); + final ArrowMessage result = new ArrowMessage(descriptor, header, appMetadata, body); + appMetadata = null; + body = null; + return result; } catch (Exception ioe) { throw new RuntimeException(ioe); + } finally { + try { + AutoCloseables.closeNoChecked(appMetadata); + } finally { + AutoCloseables.closeNoChecked(body); + } } } @@ -377,6 +403,140 @@ private static int readRawVarint32(int firstByte, InputStream is) throws IOExcep return CodedInputStream.readRawVarint32(firstByte, is); } + private static ArrowBuf readBuffer( + BufferAllocator allocator, InputStream stream, int size) throws IOException { + if (stream instanceof ArrowBufInputStream) { + return ((ArrowBufInputStream) stream).readArrowBuf(size); + } + + final ArrowBuf buffer = allocator.buffer(size); + try { + GetReadableBuffer.readIntoBuffer(stream, buffer, size, ENABLE_ZERO_COPY_READ); + return buffer; + } catch (IOException | RuntimeException | Error e) { + buffer.close(); + throw e; + } + } + + /** An InputStream over an owned gRPC buffer that can return zero-copy ArrowBuf slices. */ + private static final class ArrowBufInputStream extends InputStream { + private final ArrowBuf buffer; + private final int length; + private int position; + + private ArrowBufInputStream(ArrowBuf buffer, int length) { + this.buffer = buffer; + this.length = length; + } + + private static ArrowBufInputStream tryCreate( + BufferAllocator allocator, InputStream stream) { + if (!(stream instanceof Detachable) + || !(stream instanceof HasByteBuffer) + || !(stream instanceof KnownLength) + || !((HasByteBuffer) stream).byteBufferSupported()) { + return null; + } + + final ByteBuffer current = ((HasByteBuffer) stream).getByteBuffer(); + final int size; + try { + size = stream.available(); + } catch (IOException e) { + throw new RuntimeException("Failed to inspect gRPC input buffer", e); + } + if (current == null + || !current.isDirect() + || size == 0 + || current.remaining() != size) { + return null; + } + + final InputStream detached = ((Detachable) stream).detach(); + final ByteBuffer detachedBuffer; + final long dataAddress; + try { + if (!(detached instanceof HasByteBuffer) + || !((HasByteBuffer) detached).byteBufferSupported()) { + throw new IllegalStateException("Detached gRPC stream does not expose its ByteBuffer"); + } + detachedBuffer = ((HasByteBuffer) detached).getByteBuffer(); + if (detachedBuffer == null + || !detachedBuffer.isDirect() + || detachedBuffer.remaining() != size) { + throw new IllegalStateException("Detached gRPC input buffer changed after detaching"); + } + dataAddress = + MemoryUtil.getByteBufferAddress(detachedBuffer) + detachedBuffer.position(); + } catch (RuntimeException | Error e) { + AutoCloseables.closeNoChecked(detached); + throw e; + } + + final ArrowBuf buffer = + allocator.wrapForeignAllocation( + new ForeignAllocation(size, dataAddress) { + @Override + protected void release0() { + AutoCloseables.closeNoChecked(detached); + } + }); + return new ArrowBufInputStream(buffer, size); + } + + private ArrowBuf readArrowBuf(int size) throws IOException { + if (size < 0 || size > available()) { + throw new IOException("Unexpected end of detached gRPC input buffer"); + } + final int offset = position; + position += size; + buffer.getReferenceManager().retain(); + try { + return buffer.slice(offset, size); + } catch (RuntimeException | Error e) { + buffer.getReferenceManager().release(); + throw e; + } + } + + @Override + public int available() { + return length - position; + } + + @Override + public int read() { + return position == length ? -1 : buffer.getByte(position++) & 0xFF; + } + + @Override + public int read(byte[] bytes, int offset, int size) { + if (size == 0) { + return 0; + } + final int read = Math.min(size, available()); + if (read == 0) { + return -1; + } + buffer.getBytes(position, bytes, offset, read); + position += read; + return read; + } + + @Override + public long skip(long size) { + final int skipped = (int) Math.min(Math.max(size, 0), available()); + position += skipped; + return skipped; + } + + @Override + public void close() { + buffer.close(); + } + } + /** * Convert the ArrowMessage to an InputStream. * diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java index 15cfd6ba85..9195eb6780 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java @@ -480,6 +480,7 @@ public void onNext(ArrowMessage msg) { ex = new UnsupportedOperationException( "Unable to handle message of type: " + msg.getMessageType()); + AutoCloseables.closeNoChecked(msg); enqueue(DONE_EX); } } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index c925fbddf8..7bde4a4e3e 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -31,13 +31,10 @@ * {@code byte[]} allocation and the extra copy that goes with it. Otherwise, we fall back to * reading the stream into a heap array. * - *

This relies on gRPC's public zero-copy APIs ({@link HasByteBuffer#getByteBuffer()} and - * {@link InputStream#skip(long)}). No longer access gRPC internals via reflection. See - * apache/arrow-java#939. + *

This relies on gRPC's public buffer APIs ({@link HasByteBuffer#getByteBuffer()} and {@link + * InputStream#skip(long)}) instead of accessing gRPC internals through reflection. */ -public final class GetReadableBuffer { - - private GetReadableBuffer() {} +public class GetReadableBuffer { /** * Helper method to read a gRPC-provided InputStream into an ArrowBuf. diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java new file mode 100644 index 0000000000..4bdd0cc8df --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.common.collect.Iterables; +import com.google.protobuf.ByteString; +import com.google.protobuf.CodedOutputStream; +import com.google.protobuf.WireFormat; +import io.grpc.Detachable; +import io.grpc.HasByteBuffer; +import io.grpc.KnownLength; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import org.apache.arrow.flight.impl.Flight.FlightData; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests the detachable ArrowMessage read path. */ +public class TestArrowMessageDetachable { + private BufferAllocator allocator; + + @BeforeEach + public void setUp() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void tearDown() { + allocator.close(); + } + + @Test + public void testContiguousDirectBufferIsDetached() throws Exception { + final byte[] body = payload(64); + final byte[] serialized = serializeBody(body); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(serialized); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + assertEquals(1, stream.state.detachCount); + assertEquals(serialized.length, allocator.getAllocatedMemory()); + assertBodyEquals(message, body); + assertEquals(0, stream.state.closeCount); + } + assertEquals(1, stream.state.closeCount); + assertEquals(0, allocator.getAllocatedMemory()); + stream.close(); + } + + @Test + public void testHeapBufferFallsBackWithoutDetaching() throws Exception { + final byte[] body = payload(32); + final MockGrpcInputStream stream = MockGrpcInputStream.heap(serializeBody(body)); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + assertEquals(0, stream.state.detachCount); + assertBodyEquals(message, body); + } + assertEquals(0, allocator.getAllocatedMemory()); + stream.close(); + assertEquals(1, stream.state.closeCount); + } + + @Test + public void testDetachedBufferIsClosedOnParseFailure() throws Exception { + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + final CodedOutputStream coded = CodedOutputStream.newInstance(output); + coded.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(payload(8))); + coded.writeTag( + FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); + coded.writeUInt32NoTag(10); + coded.writeRawByte(1); + coded.flush(); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(output.toByteArray()); + + assertThrows( + RuntimeException.class, () -> ArrowMessage.createMarshaller(allocator).parse(stream)); + assertEquals(1, stream.state.detachCount); + assertEquals(1, stream.state.closeCount); + assertEquals(0, allocator.getAllocatedMemory()); + stream.close(); + } + + private static byte[] serializeBody(byte[] body) { + return FlightData.newBuilder().setDataBody(ByteString.copyFrom(body)).build().toByteArray(); + } + + private static byte[] payload(int size) { + final byte[] bytes = new byte[size]; + for (int i = 0; i < size; i++) { + bytes[i] = (byte) i; + } + return bytes; + } + + private static void assertBodyEquals(ArrowMessage message, byte[] expected) { + final ArrowBuf body = Iterables.getOnlyElement(message.getBufs()); + final byte[] actual = new byte[expected.length]; + body.getBytes(0, actual); + assertArrayEquals(expected, actual); + } + + private static final class MockGrpcInputStream extends InputStream + implements Detachable, HasByteBuffer, KnownLength { + private final State state; + private ByteBuffer buffer; + private boolean ownsBuffer = true; + private boolean closed; + + private MockGrpcInputStream(ByteBuffer buffer, State state) { + this.buffer = buffer; + this.state = state; + } + + private static MockGrpcInputStream direct(byte[] bytes) { + final ByteBuffer buffer = ByteBuffer.allocateDirect(bytes.length); + buffer.put(bytes).flip(); + return new MockGrpcInputStream(buffer, new State()); + } + + private static MockGrpcInputStream heap(byte[] bytes) { + return new MockGrpcInputStream(ByteBuffer.wrap(bytes), new State()); + } + + @Override + public boolean byteBufferSupported() { + return true; + } + + @Override + public ByteBuffer getByteBuffer() { + return buffer.hasRemaining() ? buffer.duplicate() : null; + } + + @Override + public InputStream detach() { + state.detachCount++; + final ByteBuffer detached = buffer; + buffer = ByteBuffer.allocate(0); + ownsBuffer = false; + return new MockGrpcInputStream(detached, state); + } + + @Override + public int available() { + return buffer.remaining(); + } + + @Override + public int read() { + return buffer.hasRemaining() ? buffer.get() & 0xFF : -1; + } + + @Override + public int read(byte[] bytes, int offset, int size) { + if (!buffer.hasRemaining()) { + return -1; + } + final int read = Math.min(size, buffer.remaining()); + buffer.get(bytes, offset, read); + return read; + } + + @Override + public long skip(long size) { + final int skipped = (int) Math.min(Math.max(size, 0), buffer.remaining()); + buffer.position(buffer.position() + skipped); + return skipped; + } + + @Override + public void close() throws IOException { + if (!closed) { + closed = true; + buffer = ByteBuffer.allocate(0); + if (ownsBuffer) { + ownsBuffer = false; + state.closeCount++; + } + } + } + } + + private static final class State { + private int detachCount; + private int closeCount; + } +} diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java index dd21ca699d..7b1989310f 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java @@ -73,8 +73,8 @@ public void testFastPathAcrossChunks() throws IOException { } /** - * A chunk may hold more than the caller asked for (like next field bytes). Only {@code - * size} bytes may be consumed, the rest must remain readable. + * A chunk may hold more than the caller asked for (like next field bytes). Only {@code size} + * bytes may be consumed, the rest must remain readable. */ @Test public void testFastPathDoesNotOverConsume() throws IOException { @@ -106,11 +106,12 @@ public void testFastPathSkipReturnsZero() throws IOException { public void testFastPathTruncatedStream() throws IOException { try (ChunkedStream stream = new ChunkedStream(payload(10)); ArrowBuf buf = allocator.buffer(32)) { - assertThrows(IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 32, true)); + assertThrows( + IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 32, true)); } } - /** Both take the heap-array path: streams without zero-copy support, and fastPath=false. */ + /** Both take the heap-array path: streams without ByteBuffer support, and fastPath=false. */ @Test public void testSlowPath() throws IOException { final byte[] payload = payload(33); @@ -151,8 +152,8 @@ private static byte[][] slice(byte[] payload, int... sizes) { } /** - * For gRPC's buffer-backed streams: {@link #getByteBuffer()} exposes the next chunk - * and {@link #skip(long)} advances. + * For gRPC's buffer-backed streams: {@link #getByteBuffer()} exposes the next chunk and {@link + * #skip(long)} advances. */ private static final class ChunkedStream extends InputStream implements HasByteBuffer { private final Deque chunks = new ArrayDeque<>(); From 71932aea205a153e698f3be7b4992a25379e840e Mon Sep 17 00:00:00 2001 From: mihazagar <17075688+mihazagar@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:42:59 +0200 Subject: [PATCH 4/8] Whitespace --- .../org/apache/arrow/flight/ArrowMessage.java | 15 +++++---------- .../arrow/flight/TestArrowMessageDetachable.java | 3 +-- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java index 816a380b04..440aec7f51 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java @@ -403,8 +403,8 @@ private static int readRawVarint32(int firstByte, InputStream is) throws IOExcep return CodedInputStream.readRawVarint32(firstByte, is); } - private static ArrowBuf readBuffer( - BufferAllocator allocator, InputStream stream, int size) throws IOException { + private static ArrowBuf readBuffer(BufferAllocator allocator, InputStream stream, int size) + throws IOException { if (stream instanceof ArrowBufInputStream) { return ((ArrowBufInputStream) stream).readArrowBuf(size); } @@ -430,8 +430,7 @@ private ArrowBufInputStream(ArrowBuf buffer, int length) { this.length = length; } - private static ArrowBufInputStream tryCreate( - BufferAllocator allocator, InputStream stream) { + private static ArrowBufInputStream tryCreate(BufferAllocator allocator, InputStream stream) { if (!(stream instanceof Detachable) || !(stream instanceof HasByteBuffer) || !(stream instanceof KnownLength) @@ -446,10 +445,7 @@ private static ArrowBufInputStream tryCreate( } catch (IOException e) { throw new RuntimeException("Failed to inspect gRPC input buffer", e); } - if (current == null - || !current.isDirect() - || size == 0 - || current.remaining() != size) { + if (current == null || !current.isDirect() || size == 0 || current.remaining() != size) { return null; } @@ -467,8 +463,7 @@ private static ArrowBufInputStream tryCreate( || detachedBuffer.remaining() != size) { throw new IllegalStateException("Detached gRPC input buffer changed after detaching"); } - dataAddress = - MemoryUtil.getByteBufferAddress(detachedBuffer) + detachedBuffer.position(); + dataAddress = MemoryUtil.getByteBufferAddress(detachedBuffer) + detachedBuffer.position(); } catch (RuntimeException | Error e) { AutoCloseables.closeNoChecked(detached); throw e; diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java index 4bdd0cc8df..6c9f0bb9a6 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java @@ -89,8 +89,7 @@ public void testDetachedBufferIsClosedOnParseFailure() throws Exception { final ByteArrayOutputStream output = new ByteArrayOutputStream(); final CodedOutputStream coded = CodedOutputStream.newInstance(output); coded.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(payload(8))); - coded.writeTag( - FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); + coded.writeTag(FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); coded.writeUInt32NoTag(10); coded.writeRawByte(1); coded.flush(); From fabc781a7789bb85fbb6a9806e9d7ec4c137c453 Mon Sep 17 00:00:00 2001 From: Abhishek Pathania Date: Wed, 16 Sep 2026 12:13:41 +0530 Subject: [PATCH 5/8] GH-1293: Use floor division when splitting epoch millis into day and time (#1294) ## What's Changed `DateTimeUtils.getTimestampValue(long)` used `/` and `%` to split epoch milliseconds into an epoch day and a time within that day. These operators round toward zero. For negative values that were not exactly midnight, the existing code fixed the remainder but not the epoch day. The two parts then referred to different days, so the timestamp came back one day late. For example, `-618102000000` ms is 1950-06-01 01:00:00 UTC. The old division produced epoch day `-7153`, which is 1950-06-02, while the remainder was 01:00. The method returned 1950-06-02 01:00:00. This affects DATE values before 1970 when `ArrowFlightJdbcDateVectorAccessor.getDate(Calendar)` applies a non-zero calendar offset. The offset moves the value away from midnight and exposes the division bug. Closes #1293. --- .../arrow/driver/jdbc/utils/DateTimeUtils.java | 13 +++++-------- .../arrow/driver/jdbc/utils/DateTimeUtilsTest.java | 11 +++++++++++ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtils.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtils.java index 9363e3486c..c4e7fda59b 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtils.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtils.java @@ -55,15 +55,12 @@ public static long applyCalendarOffset(long milliseconds, Calendar calendar) { * @return a {@link Timestamp} object representing the given Epoch millis */ public static Timestamp getTimestampValue(long millisWithCalendar) { - long milliseconds = millisWithCalendar; - if (milliseconds < 0) { - // LocalTime#ofNanoDay only accepts positive values - milliseconds -= ((milliseconds / MILLIS_PER_DAY) - 1) * MILLIS_PER_DAY; - } - + // Millis are negative before 1970, where only floor semantics keep the epoch day + // and the time-of-day remainder on the same day (and the remainder non-negative). return Timestamp.valueOf( LocalDateTime.of( - LocalDate.ofEpochDay(millisWithCalendar / MILLIS_PER_DAY), - LocalTime.ofNanoOfDay(TimeUnit.MILLISECONDS.toNanos(milliseconds % MILLIS_PER_DAY)))); + LocalDate.ofEpochDay(Math.floorDiv(millisWithCalendar, MILLIS_PER_DAY)), + LocalTime.ofNanoOfDay( + TimeUnit.MILLISECONDS.toNanos(Math.floorMod(millisWithCalendar, MILLIS_PER_DAY))))); } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtilsTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtilsTest.java index 9c66352023..70bcb9b4c2 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtilsTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/DateTimeUtilsTest.java @@ -95,4 +95,15 @@ public void testShouldGetTimestampNegative() { assertThat(expected, is(actual)); } + + @Test + public void testShouldGetTimestampNegativeNotAlignedToDay() { + final long epochMilli = negativeEpochMilli + 3600000L; // 1950-06-01 01:00:00 UTC + final Instant instant = Instant.ofEpochMilli(epochMilli); + + final Timestamp expected = Timestamp.from(instant); + final Timestamp actual = DateTimeUtils.getTimestampValue(epochMilli); + + assertThat(expected, is(actual)); + } } From b399919c4137269bb95039df6376ba6dfae45497 Mon Sep 17 00:00:00 2001 From: mihazagar <17075688+mihazagar@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:33:32 +0000 Subject: [PATCH 6/8] GH-939: Read gRPC message bodies through public APIs instead of reflection --- flight/flight-core/pom.xml | 1 + .../arrow/flight/grpc/GetReadableBuffer.java | 93 +++---- .../flight/grpc/TestGetReadableBuffer.java | 242 +++++++++++------- 3 files changed, 185 insertions(+), 151 deletions(-) diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml index ac05fa11f0..ee1068c85d 100644 --- a/flight/flight-core/pom.xml +++ b/flight/flight-core/pom.xml @@ -61,6 +61,7 @@ under the License. io.grpc grpc-core + runtime diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index 7bde4a4e3e..341f10b16e 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -24,27 +24,25 @@ import org.apache.arrow.memory.ArrowBuf; /** - * Copy data from a gRPC-provided InputStream into a target ArrowBuf. + * Reads a gRPC message body into an {@link ArrowBuf}. * - *

When the stream is backed by gRPC's own buffers (i.e. it implements {@link HasByteBuffer}), - * the payload is copied directly out of gRPC's {@link ByteBuffer}s, avoiding an intermediate heap - * {@code byte[]} allocation and the extra copy that goes with it. Otherwise, we fall back to - * reading the stream into a heap array. - * - *

This relies on gRPC's public buffer APIs ({@link HasByteBuffer#getByteBuffer()} and {@link - * InputStream#skip(long)}) instead of accessing gRPC internals through reflection. + *

When the stream gRPC hands us exposes its backing buffers through {@link HasByteBuffer}, the + * bytes are copied straight from those buffers into the target, skipping the intermediate heap + * array. Otherwise the stream is read into a heap array first. */ public class GetReadableBuffer { + private GetReadableBuffer() {} + /** - * Helper method to read a gRPC-provided InputStream into an ArrowBuf. + * Read exactly {@code size} bytes from {@code stream} into {@code buf}. * * @param stream The stream to read from. - * @param buf The buffer to read into. + * @param buf The buffer to read into. Its writer index is set to {@code size} on success. * @param size The number of bytes to read. - * @param fastPath Whether to enable the fast path (i.e. copy directly from the stream's backing - * {@link ByteBuffer}s when the stream supports it). - * @throws IOException if there is an error reading from the stream + * @param fastPath Whether to copy directly from the stream's backing buffers when it exposes + * them. + * @throws IOException if the stream ends early or cannot be read. */ public static void readIntoBuffer( final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath) @@ -56,56 +54,43 @@ public static void readIntoBuffer( } else { final byte[] heapBytes = new byte[size]; ByteStreams.readFully(stream, heapBytes); - buf.writeBytes(heapBytes); + buf.setBytes(0, heapBytes); } buf.writerIndex(size); } - /** - * Copy {@code size} bytes directly out of the stream's backing {@link ByteBuffer}s into {@code - * buf}. - * - *

{@link HasByteBuffer#getByteBuffer()} exposes the next chunk of readable bytes without - * advancing the stream, so we copy the chunk and then {@link InputStream#skip(long)} past (and - * release) the bytes we consumed before asking for the next chunk. - */ private static void readFromByteBuffers( - final HasByteBuffer hasByteBuffer, - final InputStream stream, - final ArrowBuf buf, - final int size) + final HasByteBuffer source, final InputStream stream, final ArrowBuf buf, final int size) throws IOException { - long writeIndex = 0; - int remaining = size; - while (remaining > 0) { - final ByteBuffer chunk = hasByteBuffer.getByteBuffer(); - // getByteBuffer() may expose more than we need (e.g. bytes belonging to the following - // field), so only copy up to the number of bytes still outstanding. We are allowed to change - // the returned buffer's position/limit without affecting the stream. - final int toRead = chunk == null ? 0 : Math.min(remaining, chunk.remaining()); - if (toRead == 0) { + int copied = 0; + while (copied < size) { + final ByteBuffer chunk = source.getByteBuffer(); + if (chunk == null || !chunk.hasRemaining()) { throw new IOException( - String.format( - "Unexpected end of stream: %d of %d bytes remaining to be read", remaining, size)); + "Unexpected end of gRPC stream: expected " + size + " bytes, got " + copied); + } + final int toRead = Math.min(size - copied, chunk.remaining()); + buf.setBytes(copied, chunk, chunk.position(), toRead); + // getByteBuffer() does not advance the stream; skip() consumes the bytes we just copied. + consume(stream, toRead); + copied += toRead; + } + } + + private static void consume(final InputStream stream, final int count) throws IOException { + int skipped = 0; + while (skipped < count) { + final long n = stream.skip(count - skipped); + if (n > 0) { + skipped += (int) n; + continue; } - chunk.limit(chunk.position() + toRead); - buf.setBytes(writeIndex, chunk); - // getByteBuffer() does not advance the stream; skip() consumes (and frees) the bytes. - long skipped = 0; - while (skipped < toRead) { - final long n = stream.skip(toRead - skipped); - if (n > 0) { - skipped += n; - } else if (stream.read() == -1) { - throw new IOException("Unexpected end of stream while consuming copied bytes"); - } else { - // InputStream.skip() is permitted to return zero without reaching EOF. We have already - // copied this byte, so read it only to guarantee that the stream makes progress. - skipped++; - } + // InputStream.skip is allowed to return 0 before the end of the stream. Force progress with + // a single-byte read, which does distinguish end of stream. + if (stream.read() == -1) { + throw new IOException("Unexpected end of gRPC stream while skipping copied bytes"); } - writeIndex += toRead; - remaining -= toRead; + skipped++; } } } diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java index 7b1989310f..589ce7df17 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java @@ -19,9 +19,10 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import io.grpc.HasByteBuffer; -import java.io.ByteArrayInputStream; +import io.grpc.KnownLength; import java.io.IOException; import java.io.InputStream; import java.nio.ByteBuffer; @@ -35,7 +36,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -/** Tests for reading a gRPC-provided {@link InputStream} into an {@link ArrowBuf}. */ +/** Tests for reading gRPC message bodies into an {@link ArrowBuf} through gRPC's public API. */ public class TestGetReadableBuffer { private BufferAllocator allocator; @@ -47,177 +48,224 @@ public void setUp() { @AfterEach public void tearDown() { + assertEquals(0, allocator.getAllocatedMemory()); allocator.close(); } @Test - public void testFastPathSingleChunk() throws IOException { + public void fastPathReadsThroughByteBuffersNotThroughHeapReads() throws IOException { final byte[] payload = payload(64); try (ChunkedStream stream = new ChunkedStream(payload); ArrowBuf buf = allocator.buffer(payload.length)) { GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); - assertEquals(payload.length, buf.writerIndex()); - assertArrayEquals(payload, toBytes(buf, payload.length)); + + assertArrayEquals(payload, contents(buf)); + assertEquals(0, stream.available()); + assertEquals(0, stream.heapReads, "fast path must not go through InputStream.read"); + assertTrue(stream.byteBufferPeeks > 0, "fast path must use HasByteBuffer.getByteBuffer"); } } - /** Copy loop stitches several backing buffers of the requested size. */ @Test - public void testFastPathAcrossChunks() throws IOException { + public void fastPathToleratesSkipReturningZero() throws IOException { + final byte[] payload = payload(48); + try (ChunkedStream stream = new ChunkedStream(true, payload); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); + + assertArrayEquals(payload, contents(buf)); + assertEquals(0, stream.available()); + } + } + + @Test + public void fastPathCopiesAcrossChunkBoundaries() throws IOException { final byte[] payload = payload(70); - try (ChunkedStream stream = new ChunkedStream(slice(payload, 10, 1, 32, 27)); + try (ChunkedStream stream = new ChunkedStream(split(payload, 10, 1, 32, 27)); ArrowBuf buf = allocator.buffer(payload.length)) { GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); - assertArrayEquals(payload, toBytes(buf, payload.length)); + + assertArrayEquals(payload, contents(buf)); + assertEquals(0, stream.available()); + assertEquals(0, stream.heapReads); } } - /** - * A chunk may hold more than the caller asked for (like next field bytes). Only {@code size} - * bytes may be consumed, the rest must remain readable. - */ @Test - public void testFastPathDoesNotOverConsume() throws IOException { - final byte[] payload = payload(48); - try (ChunkedStream stream = new ChunkedStream(payload); + public void fastPathReadsOnlyTheRequestedBytes() throws IOException { + final byte[] payload = payload(40); + try (ChunkedStream stream = new ChunkedStream(split(payload, 16, 24)); ArrowBuf buf = allocator.buffer(20)) { GetReadableBuffer.readIntoBuffer(stream, buf, 20, true); - assertArrayEquals(Arrays.copyOf(payload, 20), toBytes(buf, 20)); - final byte[] rest = new byte[payload.length - 20]; - assertEquals(rest.length, stream.read(rest)); - assertArrayEquals(Arrays.copyOfRange(payload, 20, payload.length), rest); + assertArrayEquals(Arrays.copyOf(payload, 20), contents(buf)); + assertEquals(20, stream.available()); } } - /** The copy loop must progress when a valid InputStream returns zero from skip(). */ @Test - public void testFastPathSkipReturnsZero() throws IOException { - final byte[] payload = payload(32); - try (ChunkedStream stream = new ChunkedStream(true, payload); + public void fallsBackToHeapReadWhenByteBuffersAreNotSupported() throws IOException { + final byte[] payload = payload(24); + try (ChunkedStream stream = ChunkedStream.withoutByteBufferSupport(split(payload, 9, 15)); ArrowBuf buf = allocator.buffer(payload.length)) { GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true); - assertArrayEquals(payload, toBytes(buf, payload.length)); + + assertArrayEquals(payload, contents(buf)); + assertEquals(0, stream.available()); + assertEquals(0, stream.byteBufferPeeks, "must not call getByteBuffer when unsupported"); + assertTrue(stream.heapReads > 0); } } - /** A truncated stream must fail loudly rather than leave the buffer partially filled. */ @Test - public void testFastPathTruncatedStream() throws IOException { - try (ChunkedStream stream = new ChunkedStream(payload(10)); - ArrowBuf buf = allocator.buffer(32)) { - assertThrows( - IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 32, true)); + public void fastPathDisabledUsesHeapRead() throws IOException { + final byte[] payload = payload(24); + try (ChunkedStream stream = new ChunkedStream(payload); + ArrowBuf buf = allocator.buffer(payload.length)) { + GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, false); + + assertArrayEquals(payload, contents(buf)); + assertEquals(0, stream.byteBufferPeeks, "fastPath=false must not touch HasByteBuffer"); + assertTrue(stream.heapReads > 0); } } - /** Both take the heap-array path: streams without ByteBuffer support, and fastPath=false. */ @Test - public void testSlowPath() throws IOException { - final byte[] payload = payload(33); - try (ArrowBuf buf = allocator.buffer(payload.length)) { - GetReadableBuffer.readIntoBuffer( - new ByteArrayInputStream(payload), buf, payload.length, true); - assertArrayEquals(payload, toBytes(buf, payload.length)); - } + public void fastPathThrowsWhenStreamEndsEarly() throws IOException { + final byte[] payload = payload(10); try (ChunkedStream stream = new ChunkedStream(payload); - ArrowBuf buf = allocator.buffer(payload.length)) { - GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, false); - assertArrayEquals(payload, toBytes(buf, payload.length)); + ArrowBuf buf = allocator.buffer(16)) { + assertThrows( + IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 16, true)); } } - private static byte[] payload(int size) { - final byte[] bytes = new byte[size]; - for (int i = 0; i < size; i++) { - bytes[i] = (byte) i; + private static byte[][] split(byte[] payload, int... sizes) { + final byte[][] pieces = new byte[sizes.length][]; + int offset = 0; + for (int i = 0; i < sizes.length; i++) { + pieces[i] = Arrays.copyOfRange(payload, offset, offset + sizes[i]); + offset += sizes[i]; } - return bytes; + if (offset != payload.length) { + throw new IllegalArgumentException("sizes must sum to payload length"); + } + return pieces; } - private static byte[] toBytes(ArrowBuf buf, int size) { + private static byte[] payload(int size) { final byte[] bytes = new byte[size]; - buf.getBytes(0, bytes); + for (int i = 0; i < size; i++) { + bytes[i] = (byte) (i * 7 + 3); + } return bytes; } - private static byte[][] slice(byte[] payload, int... sizes) { - final byte[][] chunks = new byte[sizes.length][]; - int offset = 0; - for (int i = 0; i < sizes.length; i++) { - chunks[i] = Arrays.copyOfRange(payload, offset, offset + sizes[i]); - offset += sizes[i]; - } - return chunks; + private static byte[] contents(ArrowBuf buf) { + final byte[] out = new byte[(int) buf.writerIndex()]; + buf.getBytes(0, out); + return out; } /** - * For gRPC's buffer-backed streams: {@link #getByteBuffer()} exposes the next chunk and {@link - * #skip(long)} advances. + * A stand-in for gRPC's message stream: a chain of direct buffers exposed through {@link + * HasByteBuffer} and {@link KnownLength}, with counters for how it was read. */ - private static final class ChunkedStream extends InputStream implements HasByteBuffer { + static final class ChunkedStream extends InputStream implements HasByteBuffer, KnownLength { private final Deque chunks = new ArrayDeque<>(); - private boolean skipReturnsZero; + private final boolean skipReturnsZeroEveryOtherCall; + private final boolean byteBufferSupported; + private boolean returnZeroFromNextSkip; + int heapReads; + int byteBufferPeeks; + + ChunkedStream(byte[]... pieces) { + this(false, true, pieces); + } - ChunkedStream(byte[]... chunks) { - this(false, chunks); + ChunkedStream(boolean skipReturnsZeroEveryOtherCall, byte[]... pieces) { + this(skipReturnsZeroEveryOtherCall, true, pieces); } - ChunkedStream(boolean skipReturnsZero, byte[]... chunks) { - this.skipReturnsZero = skipReturnsZero; - for (byte[] chunk : chunks) { - this.chunks.add(ByteBuffer.wrap(chunk)); + static ChunkedStream withoutByteBufferSupport(byte[]... pieces) { + return new ChunkedStream(false, false, pieces); + } + + private ChunkedStream( + boolean skipReturnsZeroEveryOtherCall, boolean byteBufferSupported, byte[]... pieces) { + this.skipReturnsZeroEveryOtherCall = skipReturnsZeroEveryOtherCall; + this.byteBufferSupported = byteBufferSupported; + this.returnZeroFromNextSkip = skipReturnsZeroEveryOtherCall; + for (byte[] piece : pieces) { + final ByteBuffer chunk = ByteBuffer.allocateDirect(piece.length); + chunk.put(piece).flip(); + chunks.add(chunk); + } + } + + private ByteBuffer current() { + while (!chunks.isEmpty() && !chunks.peek().hasRemaining()) { + chunks.poll(); } + return chunks.peek(); } @Override - public boolean byteBufferSupported() { - return true; + public int read() { + final ByteBuffer chunk = current(); + return chunk == null ? -1 : chunk.get() & 0xFF; } @Override - public ByteBuffer getByteBuffer() { - final ByteBuffer head = chunks.peek(); - // Like gRPC, hand out an independent view so the caller may adjust position/limit freely. - return head == null ? null : head.duplicate(); + public int read(byte[] b, int off, int len) { + heapReads++; + final ByteBuffer chunk = current(); + if (chunk == null) { + return -1; + } + final int n = Math.min(len, chunk.remaining()); + chunk.get(b, off, n); + return n; } @Override public long skip(long n) { - if (skipReturnsZero) { - skipReturnsZero = false; - return 0; + if (skipReturnsZeroEveryOtherCall) { + returnZeroFromNextSkip = !returnZeroFromNextSkip; + if (!returnZeroFromNextSkip) { + // InputStream.skip may legitimately return 0 before end of stream. + return 0; + } } - final ByteBuffer head = chunks.peek(); - if (head == null) { + final ByteBuffer chunk = current(); + if (chunk == null) { return 0; } - final int skipped = (int) Math.min(n, head.remaining()); - head.position(head.position() + skipped); - if (!head.hasRemaining()) { - chunks.poll(); - } + // Like gRPC, skip at most one chunk per call. + final int skipped = (int) Math.min(n, chunk.remaining()); + chunk.position(chunk.position() + skipped); return skipped; } @Override - public int read() { - final byte[] one = new byte[1]; - return read(one, 0, 1) == 1 ? one[0] & 0xFF : -1; + public int available() { + int total = 0; + for (ByteBuffer chunk : chunks) { + total += chunk.remaining(); + } + return total; } @Override - public int read(byte[] dst, int off, int len) { - final ByteBuffer head = chunks.peek(); - if (head == null) { - return -1; - } - final int read = Math.min(len, head.remaining()); - head.get(dst, off, read); - if (!head.hasRemaining()) { - chunks.poll(); - } - return read; + public boolean byteBufferSupported() { + return byteBufferSupported; + } + + @Override + public ByteBuffer getByteBuffer() { + byteBufferPeeks++; + final ByteBuffer chunk = current(); + return chunk == null ? null : chunk.duplicate(); } } } From 9f9291fc6e0e3a6c2e79fe760a59d2c4a6e99d38 Mon Sep 17 00:00:00 2001 From: mihazagar <17075688+mihazagar@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:40:58 +0000 Subject: [PATCH 7/8] GH-939: Parse FlightData in place when gRPC lets us own its buffer --- .../org/apache/arrow/flight/ArrowMessage.java | 176 +++++++------ .../flight/TestArrowMessageDetachable.java | 240 ++++++++++++------ 2 files changed, 252 insertions(+), 164 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java index 440aec7f51..325671fc94 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java @@ -286,21 +286,23 @@ public Iterable getBufs() { private static ArrowMessage frame(BufferAllocator allocator, final InputStream stream) { if (ENABLE_ZERO_COPY_READ) { - final ArrowBufInputStream detached = ArrowBufInputStream.tryCreate(allocator, stream); - if (detached != null) { + // If gRPC lets us take ownership of the whole message, parse it in place: the same loop + // below then runs over our own buffer and hands out slices instead of copies. + final OwnedFrame owned = OwnedFrame.tryTakeOwnership(allocator, stream); + if (owned != null) { try { - return frame(allocator, detached); + return frame(allocator, owned); } finally { - detached.close(); + owned.close(); } } } - FlightDescriptor descriptor = null; - MessageMetadataResult header = null; ArrowBuf body = null; ArrowBuf appMetadata = null; try { + FlightDescriptor descriptor = null; + MessageMetadataResult header = null; while (stream.available() > 0) { final int tagFirstByte = stream.read(); if (tagFirstByte == -1) { @@ -326,10 +328,6 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s } case APP_METADATA_TAG: { - if (appMetadata != null) { - appMetadata.close(); - appMetadata = null; - } int size = readRawVarint32(stream); appMetadata = readBuffer(allocator, stream, size); break; @@ -337,7 +335,7 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s case BODY_TAG: if (body != null) { // only read last body. - body.close(); + body.getReferenceManager().release(); body = null; } int size = readRawVarint32(stream); @@ -379,18 +377,17 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s break; } } - final ArrowMessage result = new ArrowMessage(descriptor, header, appMetadata, body); + final ArrowMessage message = new ArrowMessage(descriptor, header, appMetadata, body); + // Ownership of the buffers has moved into the message. appMetadata = null; body = null; - return result; + return message; } catch (Exception ioe) { throw new RuntimeException(ioe); } finally { - try { - AutoCloseables.closeNoChecked(appMetadata); - } finally { - AutoCloseables.closeNoChecked(body); - } + // Only non-null if parsing failed part-way through. + AutoCloseables.closeNoChecked(appMetadata); + AutoCloseables.closeNoChecked(body); } } @@ -399,16 +396,12 @@ private static int readRawVarint32(InputStream is) throws IOException { return readRawVarint32(firstByte, is); } - private static int readRawVarint32(int firstByte, InputStream is) throws IOException { - return CodedInputStream.readRawVarint32(firstByte, is); - } - + /** Read {@code size} bytes of a length-delimited field into a buffer the caller must close. */ private static ArrowBuf readBuffer(BufferAllocator allocator, InputStream stream, int size) throws IOException { - if (stream instanceof ArrowBufInputStream) { - return ((ArrowBufInputStream) stream).readArrowBuf(size); + if (stream instanceof OwnedFrame) { + return ((OwnedFrame) stream).slice(size); } - final ArrowBuf buffer = allocator.buffer(size); try { GetReadableBuffer.readIntoBuffer(stream, buffer, size, ENABLE_ZERO_COPY_READ); @@ -419,119 +412,142 @@ private static ArrowBuf readBuffer(BufferAllocator allocator, InputStream stream } } - /** An InputStream over an owned gRPC buffer that can return zero-copy ArrowBuf slices. */ - private static final class ArrowBufInputStream extends InputStream { - private final ArrowBuf buffer; + /** + * A gRPC message we own outright, readable as a stream and sliceable without copying. + * + *

gRPC closes the stream it passes to a marshaller as soon as parsing returns, so a buffer + * merely borrowed through {@link HasByteBuffer} cannot outlive {@code parse()}. {@link + * Detachable#detach()} transfers ownership of the buffer to us; we wrap that memory as a foreign + * allocation whose release closes the detached stream, and hand out retained slices of it. The + * memory goes back to gRPC once the message and every slice are closed. + */ + private static final class OwnedFrame extends InputStream { + private final ArrowBuf frame; private final int length; private int position; - private ArrowBufInputStream(ArrowBuf buffer, int length) { - this.buffer = buffer; + private OwnedFrame(ArrowBuf frame, int length) { + this.frame = frame; this.length = length; } - private static ArrowBufInputStream tryCreate(BufferAllocator allocator, InputStream stream) { + /** + * Take ownership of the stream's buffer if gRPC exposes it and it holds the entire message. + * + * @return the owned frame, or null to fall back to copying. + */ + static OwnedFrame tryTakeOwnership(BufferAllocator allocator, InputStream stream) { if (!(stream instanceof Detachable) || !(stream instanceof HasByteBuffer) || !(stream instanceof KnownLength) || !((HasByteBuffer) stream).byteBufferSupported()) { return null; } - - final ByteBuffer current = ((HasByteBuffer) stream).getByteBuffer(); final int size; try { size = stream.available(); } catch (IOException e) { - throw new RuntimeException("Failed to inspect gRPC input buffer", e); + throw new RuntimeException("Failed to query gRPC stream length", e); } - if (current == null || !current.isDirect() || size == 0 || current.remaining() != size) { + final ByteBuffer peek = ((HasByteBuffer) stream).getByteBuffer(); + if (peek == null || !peek.isDirect() || size == 0 || peek.remaining() != size) { + // Empty, on-heap, or split across several buffers: only a single direct buffer holding + // the whole message can be wrapped. return null; } - final InputStream detached = ((Detachable) stream).detach(); - final ByteBuffer detachedBuffer; - final long dataAddress; + final InputStream owner = ((Detachable) stream).detach(); + final long address; try { - if (!(detached instanceof HasByteBuffer) - || !((HasByteBuffer) detached).byteBufferSupported()) { - throw new IllegalStateException("Detached gRPC stream does not expose its ByteBuffer"); + if (!(owner instanceof HasByteBuffer) || !((HasByteBuffer) owner).byteBufferSupported()) { + throw new IllegalStateException("Detached gRPC stream does not expose its buffer"); } - detachedBuffer = ((HasByteBuffer) detached).getByteBuffer(); - if (detachedBuffer == null - || !detachedBuffer.isDirect() - || detachedBuffer.remaining() != size) { - throw new IllegalStateException("Detached gRPC input buffer changed after detaching"); + final ByteBuffer owned = ((HasByteBuffer) owner).getByteBuffer(); + if (owned == null || !owned.isDirect() || owned.remaining() != size) { + throw new IllegalStateException("Detached gRPC buffer differs from the one inspected"); } - dataAddress = MemoryUtil.getByteBufferAddress(detachedBuffer) + detachedBuffer.position(); + address = MemoryUtil.getByteBufferAddress(owned) + owned.position(); } catch (RuntimeException | Error e) { - AutoCloseables.closeNoChecked(detached); + AutoCloseables.closeNoChecked(owner); throw e; } - final ArrowBuf buffer = - allocator.wrapForeignAllocation( - new ForeignAllocation(size, dataAddress) { - @Override - protected void release0() { - AutoCloseables.closeNoChecked(detached); - } - }); - return new ArrowBufInputStream(buffer, size); + final ArrowBuf frame; + try { + frame = + allocator.wrapForeignAllocation( + new ForeignAllocation(size, address) { + @Override + protected void release0() { + AutoCloseables.closeNoChecked(owner); + } + }); + } catch (RuntimeException | Error e) { + // wrapForeignAllocation failed before adopting the allocation, so release0 will not run. + AutoCloseables.closeNoChecked(owner); + throw e; + } + return new OwnedFrame(frame, size); } - private ArrowBuf readArrowBuf(int size) throws IOException { + /** A zero-copy view of the next {@code size} bytes; the caller owns the returned buffer. */ + ArrowBuf slice(int size) throws IOException { if (size < 0 || size > available()) { - throw new IOException("Unexpected end of detached gRPC input buffer"); + throw new IOException("Unexpected end of gRPC message"); } final int offset = position; position += size; - buffer.getReferenceManager().retain(); + // A slice shares the frame's reference count, so give it its own reference first. + frame.getReferenceManager().retain(); try { - return buffer.slice(offset, size); + return frame.slice(offset, size); } catch (RuntimeException | Error e) { - buffer.getReferenceManager().release(); + frame.getReferenceManager().release(); throw e; } } - @Override - public int available() { - return length - position; - } - @Override public int read() { - return position == length ? -1 : buffer.getByte(position++) & 0xFF; + if (position >= length) { + return -1; + } + return frame.getByte(position++) & 0xFF; } @Override - public int read(byte[] bytes, int offset, int size) { - if (size == 0) { - return 0; - } - final int read = Math.min(size, available()); - if (read == 0) { + public int read(byte[] b, int off, int len) { + if (position >= length) { return -1; } - buffer.getBytes(position, bytes, offset, read); - position += read; - return read; + final int n = Math.min(len, available()); + frame.getBytes(position, b, off, n); + position += n; + return n; } @Override - public long skip(long size) { - final int skipped = (int) Math.min(Math.max(size, 0), available()); + public long skip(long n) { + final int skipped = (int) Math.min(n, available()); position += skipped; return skipped; } + @Override + public int available() { + return length - position; + } + @Override public void close() { - buffer.close(); + frame.close(); } } + private static int readRawVarint32(int firstByte, InputStream is) throws IOException { + return CodedInputStream.readRawVarint32(firstByte, is); + } + /** * Convert the ArrowMessage to an InputStream. * diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java index 6c9f0bb9a6..0d4661723f 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java @@ -20,7 +20,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; -import com.google.common.collect.Iterables; import com.google.protobuf.ByteString; import com.google.protobuf.CodedOutputStream; import com.google.protobuf.WireFormat; @@ -28,7 +27,6 @@ import io.grpc.HasByteBuffer; import io.grpc.KnownLength; import java.io.ByteArrayOutputStream; -import java.io.IOException; import java.io.InputStream; import java.nio.ByteBuffer; import org.apache.arrow.flight.impl.Flight.FlightData; @@ -39,8 +37,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -/** Tests the detachable ArrowMessage read path. */ +/** Tests for parsing FlightData without copying, by taking ownership of gRPC's buffer. */ public class TestArrowMessageDetachable { + private BufferAllocator allocator; @BeforeEach @@ -50,117 +49,196 @@ public void setUp() { @AfterEach public void tearDown() { + assertEquals(0, allocator.getAllocatedMemory()); allocator.close(); } @Test - public void testContiguousDirectBufferIsDetached() throws Exception { + public void contiguousDirectBufferIsDetachedAndWrappedWithoutCopy() throws Exception { + final byte[] metadata = payload(16); final byte[] body = payload(64); - final byte[] serialized = serializeBody(body); + final byte[] serialized = flightData(metadata, body); final MockGrpcInputStream stream = MockGrpcInputStream.direct(serialized); try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { - assertEquals(1, stream.state.detachCount); - assertEquals(serialized.length, allocator.getAllocatedMemory()); - assertBodyEquals(message, body); - assertEquals(0, stream.state.closeCount); + // gRPC closes the stream it handed us as soon as parse() returns. + stream.close(); + + assertEquals(1, stream.state.detachCount, "must take ownership through detach()"); + assertEquals( + serialized.length, + allocator.getAllocatedMemory(), + "the whole gRPC frame is wrapped, not a copy of the body"); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); + assertEquals( + 0, stream.state.detachedCloseCount, "gRPC memory must stay alive with the message"); } - assertEquals(1, stream.state.closeCount); - assertEquals(0, allocator.getAllocatedMemory()); - stream.close(); + assertEquals( + 1, stream.state.detachedCloseCount, "closing the message releases gRPC's buffer once"); } @Test - public void testHeapBufferFallsBackWithoutDetaching() throws Exception { - final byte[] body = payload(32); - final MockGrpcInputStream stream = MockGrpcInputStream.heap(serializeBody(body)); + public void heapBufferFallsBackToCopyingWithoutDetaching() throws Exception { + final byte[] metadata = payload(16); + final byte[] body = payload(64); + final MockGrpcInputStream stream = MockGrpcInputStream.heap(flightData(metadata, body)); try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { - assertEquals(0, stream.state.detachCount); - assertBodyEquals(message, body); + stream.close(); + + assertEquals(0, stream.state.detachCount, "on-heap buffers cannot be wrapped"); + assertEquals( + metadata.length + body.length, + allocator.getAllocatedMemory(), + "copying path allocates exactly the two fields"); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); } - assertEquals(0, allocator.getAllocatedMemory()); - stream.close(); - assertEquals(1, stream.state.closeCount); + assertEquals(0, stream.state.detachedCloseCount); } @Test - public void testDetachedBufferIsClosedOnParseFailure() throws Exception { - final ByteArrayOutputStream output = new ByteArrayOutputStream(); - final CodedOutputStream coded = CodedOutputStream.newInstance(output); - coded.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(payload(8))); - coded.writeTag(FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); - coded.writeUInt32NoTag(10); - coded.writeRawByte(1); - coded.flush(); - final MockGrpcInputStream stream = MockGrpcInputStream.direct(output.toByteArray()); + public void messageSplitAcrossBuffersFallsBackToCopying() throws Exception { + final byte[] metadata = payload(16); + final byte[] body = payload(64); + final byte[] serialized = flightData(metadata, body); + final MockGrpcInputStream stream = MockGrpcInputStream.directFragmented(serialized, 20); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + stream.close(); + + assertEquals(0, stream.state.detachCount, "only a single contiguous buffer is wrapped"); + assertEquals(metadata.length + body.length, allocator.getAllocatedMemory()); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); + } + } + + @Test + public void parseFailureReleasesTheDetachedBuffer() throws Exception { + // A complete body, then a descriptor whose declared length runs past the end of the frame. + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(bytes); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(payload(8))); + out.writeTag(FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); + out.writeUInt32NoTag(10); + out.writeRawByte(1); + out.flush(); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(bytes.toByteArray()); assertThrows( RuntimeException.class, () -> ArrowMessage.createMarshaller(allocator).parse(stream)); + stream.close(); + assertEquals(1, stream.state.detachCount); - assertEquals(1, stream.state.closeCount); + assertEquals(1, stream.state.detachedCloseCount, "the detached buffer must not leak"); assertEquals(0, allocator.getAllocatedMemory()); - stream.close(); } - private static byte[] serializeBody(byte[] body) { - return FlightData.newBuilder().setDataBody(ByteString.copyFrom(body)).build().toByteArray(); + @Test + public void lastOfDuplicateBodyFieldsWinsWithoutLeaking() throws Exception { + final byte[] first = payload(8); + final byte[] second = payload(24); + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(bytes); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(first)); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(second)); + out.flush(); + final byte[] serialized = bytes.toByteArray(); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(serialized); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + stream.close(); + assertEquals(1, stream.state.detachCount); + assertArrayEquals(second, contents(message.getBufs().iterator().next())); + assertEquals(serialized.length, allocator.getAllocatedMemory()); + } + assertEquals(1, stream.state.detachedCloseCount); + } + + private static byte[] flightData(byte[] metadata, byte[] body) { + return FlightData.newBuilder() + .setAppMetadata(ByteString.copyFrom(metadata)) + .setDataBody(ByteString.copyFrom(body)) + .build() + .toByteArray(); } private static byte[] payload(int size) { final byte[] bytes = new byte[size]; for (int i = 0; i < size; i++) { - bytes[i] = (byte) i; + bytes[i] = (byte) (i * 13 + 7); } return bytes; } - private static void assertBodyEquals(ArrowMessage message, byte[] expected) { - final ArrowBuf body = Iterables.getOnlyElement(message.getBufs()); - final byte[] actual = new byte[expected.length]; - body.getBytes(0, actual); - assertArrayEquals(expected, actual); + private static byte[] contents(ArrowBuf buf) { + final byte[] out = new byte[(int) buf.writerIndex()]; + buf.getBytes(0, out); + return out; } - private static final class MockGrpcInputStream extends InputStream + /** + * A stand-in for the stream gRPC's Netty transport hands to a marshaller: one buffer, exposed + * through the three public capabilities, with counters for detach and close. + */ + static final class MockGrpcInputStream extends InputStream implements Detachable, HasByteBuffer, KnownLength { - private final State state; + + static final class State { + int detachCount; + int detachedCloseCount; + } + + final State state; private ByteBuffer buffer; - private boolean ownsBuffer = true; - private boolean closed; + private final boolean detached; - private MockGrpcInputStream(ByteBuffer buffer, State state) { - this.buffer = buffer; - this.state = state; + /** How many bytes getByteBuffer() exposes at once; gRPC may hold a message in pieces. */ + private final int exposeLimit; + + static MockGrpcInputStream direct(byte[] bytes) { + return directFragmented(bytes, Integer.MAX_VALUE); } - private static MockGrpcInputStream direct(byte[] bytes) { + static MockGrpcInputStream directFragmented(byte[] bytes, int exposeLimit) { final ByteBuffer buffer = ByteBuffer.allocateDirect(bytes.length); buffer.put(bytes).flip(); - return new MockGrpcInputStream(buffer, new State()); + return new MockGrpcInputStream(buffer, new State(), false, exposeLimit); } - private static MockGrpcInputStream heap(byte[] bytes) { - return new MockGrpcInputStream(ByteBuffer.wrap(bytes), new State()); + static MockGrpcInputStream heap(byte[] bytes) { + return new MockGrpcInputStream(ByteBuffer.wrap(bytes), new State(), false, Integer.MAX_VALUE); + } + + private MockGrpcInputStream(ByteBuffer buffer, State state, boolean detached, int exposeLimit) { + this.buffer = buffer; + this.state = state; + this.detached = detached; + this.exposeLimit = exposeLimit; } @Override - public boolean byteBufferSupported() { - return true; + public int read() { + return buffer.hasRemaining() ? buffer.get() & 0xFF : -1; } @Override - public ByteBuffer getByteBuffer() { - return buffer.hasRemaining() ? buffer.duplicate() : null; + public int read(byte[] b, int off, int len) { + if (!buffer.hasRemaining()) { + return -1; + } + final int n = Math.min(len, buffer.remaining()); + buffer.get(b, off, n); + return n; } @Override - public InputStream detach() { - state.detachCount++; - final ByteBuffer detached = buffer; - buffer = ByteBuffer.allocate(0); - ownsBuffer = false; - return new MockGrpcInputStream(detached, state); + public long skip(long n) { + final int skipped = (int) Math.min(n, buffer.remaining()); + buffer.position(buffer.position() + skipped); + return skipped; } @Override @@ -169,42 +247,36 @@ public int available() { } @Override - public int read() { - return buffer.hasRemaining() ? buffer.get() & 0xFF : -1; + public boolean byteBufferSupported() { + return true; } @Override - public int read(byte[] bytes, int offset, int size) { + public ByteBuffer getByteBuffer() { if (!buffer.hasRemaining()) { - return -1; + return null; + } + final ByteBuffer view = buffer.duplicate(); + if (exposeLimit < view.remaining()) { + view.limit(view.position() + exposeLimit); } - final int read = Math.min(size, buffer.remaining()); - buffer.get(bytes, offset, read); - return read; + return view; } @Override - public long skip(long size) { - final int skipped = (int) Math.min(Math.max(size, 0), buffer.remaining()); - buffer.position(buffer.position() + skipped); - return skipped; + public InputStream detach() { + state.detachCount++; + final MockGrpcInputStream owner = new MockGrpcInputStream(buffer, state, true, exposeLimit); + buffer = ByteBuffer.allocate(0); + return owner; } @Override - public void close() throws IOException { - if (!closed) { - closed = true; - buffer = ByteBuffer.allocate(0); - if (ownsBuffer) { - ownsBuffer = false; - state.closeCount++; - } + public void close() { + if (detached) { + state.detachedCloseCount++; } + buffer = ByteBuffer.allocate(0); } } - - private static final class State { - private int detachCount; - private int closeCount; - } } From 0e31a854feaf81e4f5cdee1e31aec26d623dcfef Mon Sep 17 00:00:00 2001 From: mihazagar <17075688+mihazagar@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:58:38 +0000 Subject: [PATCH 8/8] GH-939: Prove the zero-copy read path fires over a real gRPC transport --- .../arrow/flight/TestFlightZeroCopyRead.java | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java new file mode 100644 index 0000000000..4431f301df --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; +import static org.apache.arrow.flight.Location.forGrpcInsecure; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.common.io.ByteStreams; +import io.grpc.Drainable; +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.message.IpcOption; +import org.junit.jupiter.api.Test; + +/** + * Proves the zero-copy read path is taken over a real gRPC transport, not only against a mock. + * + *

The oracle is allocator accounting. When the client takes ownership of gRPC's buffer, it is + * charged for the whole wire frame (header, tags and body). When it copies, it is charged for a + * buffer of the body's size only. The two are never equal, so the number distinguishes the paths. + */ +public class TestFlightZeroCopyRead { + + private static final int ROWS = 1000; + + @Test + public void recordBatchIsReadWithoutCopyingOverRealTransport() throws Exception { + assertTrue(ArrowMessage.ENABLE_ZERO_COPY_READ, "zero-copy reads must be enabled for this test"); + try (BufferAllocator serverAllocator = new RootAllocator(Long.MAX_VALUE); + BufferAllocator clientAllocator = new RootAllocator(Long.MAX_VALUE); + FlightServer server = + FlightServer.builder( + serverAllocator, + forGrpcInsecure(LOCALHOST, 0), + new OneBatchProducer(serverAllocator)) + .build() + .start(); + FlightClient client = FlightClient.builder(clientAllocator, server.getLocation()).build()) { + final long wireSize = wireSizeOfTheBatch(clientAllocator); + assertEquals(0, clientAllocator.getAllocatedMemory()); + + try (FlightStream stream = client.getStream(new Ticket(new byte[0]))) { + assertTrue(stream.next()); + final VectorSchemaRoot root = stream.getRoot(); + final IntVector values = (IntVector) root.getVector("c1"); + assertEquals(ROWS, root.getRowCount()); + for (int i = 0; i < ROWS; i++) { + assertEquals(i * 3, values.get(i)); + } + assertEquals( + wireSize, + clientAllocator.getAllocatedMemory(), + "the client should be holding gRPC's whole frame, not a copy of the body"); + assertFalse(stream.next()); + } + assertEquals(0, clientAllocator.getAllocatedMemory(), "gRPC's buffer must be released"); + } + } + + /** The exact number of bytes the server puts on the wire for the batch the producer sends. */ + private static long wireSizeOfTheBatch(BufferAllocator allocator) throws Exception { + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + // Same ownership convention as the server: the batch's buffer references belong to the + // message, which is closed once the wire stream is done with them. + try (VectorSchemaRoot root = theBatch(allocator); + ArrowMessage message = + new ArrowMessage( + new VectorUnloader(root).getRecordBatch(), null, false, IpcOption.DEFAULT); + InputStream wire = ArrowMessage.createMarshaller(allocator).stream(message)) { + if (wire instanceof Drainable) { + ((Drainable) wire).drainTo(bytes); + } else { + ByteStreams.copy(wire, bytes); + } + } + return bytes.size(); + } + + private static VectorSchemaRoot theBatch(BufferAllocator allocator) { + final IntVector values = new IntVector("c1", allocator); + final VectorSchemaRoot root = VectorSchemaRoot.of(values); + root.allocateNew(); + for (int i = 0; i < ROWS; i++) { + values.set(i, i * 3); + } + values.setValueCount(ROWS); + root.setRowCount(ROWS); + return root; + } + + private static final class OneBatchProducer extends NoOpFlightProducer { + private final BufferAllocator allocator; + + OneBatchProducer(BufferAllocator allocator) { + this.allocator = allocator; + } + + @Override + public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { + try (VectorSchemaRoot root = theBatch(allocator)) { + listener.start(root); + listener.putNext(); + listener.completed(); + } + } + } +}