fix: Yield buffered bytes before pulling the next chunk in stream adapters

src/io/adapters.ts

@@ -93,7 +93,7 @@ function* fromIterable<T extends ArrayBufferViewInput>(source: Iterable<T> | T):
             if (done || size <= bufferLength) {
                 do {
                     ({ cmd, size } = yield byteRange());
-                } while (size < bufferLength);
+                } while (size < bufferLength || (size > 0 && size === bufferLength));
             }
         } while (!done);
     } catch (e) {
@@ -141,7 +141,7 @@ async function* fromAsyncIterable<T extends ArrayBufferViewInput>(source: AsyncI
             if (done || size <= bufferLength) {
                 do {
                     ({ cmd, size } = yield byteRange());
-                } while (size < bufferLength);
+                } while (size < bufferLength || (size > 0 && size === bufferLength));
             }
         } while (!done);
     } catch (e) {
@@ -192,7 +192,7 @@ async function* fromDOMStream<T extends ArrayBufferViewInput>(source: ReadableSt
             if (done || size <= bufferLength) {
                 do {
                     ({ cmd, size } = yield byteRange());
-                } while (size < bufferLength);
+                } while (size < bufferLength || (size > 0 && size === bufferLength));
             }
         } while (!done);
     } catch (e) {
@@ -324,7 +324,7 @@ async function* fromNodeStream(stream: NodeJS.ReadableStream): AsyncUint8ArrayGe
             if (done || size <= bufferLength) {
                 do {
                     ({ cmd, size } = yield byteRange());
-                } while (size < bufferLength);
+                } while (size < bufferLength || (size > 0 && size === bufferLength));
             }
         } while (!done);
     } finally {

test/unit/ipc/reader/stream-reader-tests.ts

@@ -22,7 +22,15 @@ import {
     validateRecordBatchReader
 } from '../validate.js';
 
-import { RecordBatchReader } from 'apache-arrow';
+import {
+    AsyncByteStream,
+    ByteStream,
+    Message,
+    RecordBatchReader,
+    RecordBatchStreamWriter,
+    Table,
+    tableFromArrays
+} from 'apache-arrow';
 
 for (const table of generateRandomTables([10, 20, 30])) {
 
@@ -59,3 +67,84 @@ for (const table of generateRandomTables([10, 20, 30])) {
         });
     });
 }
+
+describe(`RecordBatchStreamReader over message-sized chunks`, () => {
+
+    const source = tableFromArrays({ x: [0, 10, 20, 30, 40, 50] });
+    const batch = source.batches[0];
+    const table = new Table(source.schema, [batch.slice(0, 2), batch.slice(2, 4), batch.slice(4, 6)]);
+
+    test(`AsyncIterable yields each RecordBatch as soon as its chunk arrives`, async () => {
+        const messages = [...ipcMessages(await RecordBatchStreamWriter.writeAll(table).toUint8Array())];
+        let pulled = 0;
+        async function* chunks() {
+            for (const message of messages) { pulled++; yield message; }
+            // the byte stream stays open, like a WebSocket or network stream
+            await new Promise(() => { });
+        }
+
+        const reader = await (await RecordBatchReader.from(chunks())).open();
+        const iterator = reader[Symbol.asyncIterator]();
+
+        // reading the schema should not pull the next chunk
+        expect(pulled).toBe(1);
+        for (let i = 0; i < table.batches.length; i++) {
+            const { value } = await iterator.next();
+            expect(value).toEqualRecordBatch(table.batches[i]);
+            expect(pulled).toBe(i + 2);
+        }
+        await iterator.return!(null);
+    });
+
+    test(`AsyncIterable terminates when the byte stream ends`, async () => {
+        const messages = [...ipcMessages(await RecordBatchStreamWriter.writeAll(table).toUint8Array())];
+        const reader = await (await RecordBatchReader.from(async function* () { yield* messages; }())).open();
+        let i = 0;
+        for await (const batch of reader) {
+            expect(batch).toEqualRecordBatch(table.batches[i++]);
+        }
+        expect(i).toBe(table.batches.length);
+    });
+
+    test(`Iterable yields each RecordBatch as soon as its chunk arrives`, () => {
+        const messages = [...ipcMessages(RecordBatchStreamWriter.writeAll(table).toUint8Array(true))];
+        let pulled = 0;
+        function* chunks() {
+            for (const message of messages) { pulled++; yield message; }
+        }
+
+        const reader = RecordBatchReader.from(chunks()).open();
+        expect(pulled).toBe(1);
+
+        let i = 0;
+        for (const batch of reader) {
+            expect(batch).toEqualRecordBatch(table.batches[i]);
+            expect(pulled).toBe(++i + 1);
+        }
+        expect(i).toBe(table.batches.length);
+    });
+});
+
+test.each([null, 0])(`ByteStream and AsyncByteStream advance after read(%s)`, async (size) => {
+    const chunks = [new Uint8Array([1]), new Uint8Array([2])];
+    for (const stream of [new ByteStream(chunks), new AsyncByteStream(async function* () { yield* chunks; }())]) {
+        for (const chunk of chunks) {
+            expect(await stream.read(size)).toEqual(chunk);
+        }
+        expect(await stream.read(size)).toEqual(new Uint8Array(0));
+        expect((await stream.next(size)).done).toBe(true);
+    }
+});
+
+// Splits an IPC stream into one Uint8Array per encapsulated message
+// (4- or 8-byte prefix, flatbuffer metadata + padding, body, EOS).
+function* ipcMessages(buffer: Uint8Array) {
+    const view = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength);
+    for (let offset = 0; offset < buffer.byteLength;) {
+        const prefixSize = view.getInt32(offset, true) === -1 ? 8 : 4;
+        const metadataLength = view.getInt32(offset + (prefixSize - 4), true);
+        const bodyLength = metadataLength === 0 ? 0 : Number(
+            Message.decode(buffer.subarray(offset + prefixSize, offset + prefixSize + metadataLength)).bodyLength);
+        yield buffer.subarray(offset, offset += prefixSize + metadataLength + bodyLength);
+    }
+}