From 141e78bfa4a7cf454ad5c2c74b9e0b41e545a79e Mon Sep 17 00:00:00 2001 From: Christopher Pruijsen Date: Tue, 13 Oct 2026 19:03:41 +0100 Subject: [PATCH 1/2] fix: Yield buffered bytes before pulling the next chunk in stream adapters --- src/io/adapters.ts | 8 +-- test/unit/ipc/reader/stream-reader-tests.ts | 78 ++++++++++++++++++++- 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/src/io/adapters.ts b/src/io/adapters.ts index 7d25e001..e970282c 100644 --- a/src/io/adapters.ts +++ b/src/io/adapters.ts @@ -93,7 +93,7 @@ function* fromIterable(source: Iterable | T): if (done || size <= bufferLength) { do { ({ cmd, size } = yield byteRange()); - } while (size < bufferLength); + } while (size <= bufferLength); } } while (!done); } catch (e) { @@ -141,7 +141,7 @@ async function* fromAsyncIterable(source: AsyncI if (done || size <= bufferLength) { do { ({ cmd, size } = yield byteRange()); - } while (size < bufferLength); + } while (size <= bufferLength); } } while (!done); } catch (e) { @@ -192,7 +192,7 @@ async function* fromDOMStream(source: ReadableSt if (done || size <= bufferLength) { do { ({ cmd, size } = yield byteRange()); - } while (size < bufferLength); + } while (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); } } while (!done); } finally { diff --git a/test/unit/ipc/reader/stream-reader-tests.ts b/test/unit/ipc/reader/stream-reader-tests.ts index 0ce596e1..d0022840 100644 --- a/test/unit/ipc/reader/stream-reader-tests.ts +++ b/test/unit/ipc/reader/stream-reader-tests.ts @@ -22,7 +22,13 @@ import { validateRecordBatchReader } from '../validate.js'; -import { RecordBatchReader } from 'apache-arrow'; +import { + Message, + RecordBatchReader, + RecordBatchStreamWriter, + Table, + tableFromArrays +} from 'apache-arrow'; for (const table of generateRandomTables([10, 20, 30])) { @@ -59,3 +65,73 @@ 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); + }); +}); + +// 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); + } +} From a84643562e38b375dc54dbfde96fe19673ee13d4 Mon Sep 17 00:00:00 2001 From: Christopher Pruijsen Date: Tue, 13 Oct 2026 19:12:50 +0100 Subject: [PATCH 2/2] fix: Preserve unbounded byte stream reads --- src/io/adapters.ts | 8 ++++---- test/unit/ipc/reader/stream-reader-tests.ts | 13 +++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/io/adapters.ts b/src/io/adapters.ts index e970282c..cd4e1956 100644 --- a/src/io/adapters.ts +++ b/src/io/adapters.ts @@ -93,7 +93,7 @@ function* fromIterable(source: Iterable | 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(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(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 { diff --git a/test/unit/ipc/reader/stream-reader-tests.ts b/test/unit/ipc/reader/stream-reader-tests.ts index d0022840..0cf37838 100644 --- a/test/unit/ipc/reader/stream-reader-tests.ts +++ b/test/unit/ipc/reader/stream-reader-tests.ts @@ -23,6 +23,8 @@ import { } from '../validate.js'; import { + AsyncByteStream, + ByteStream, Message, RecordBatchReader, RecordBatchStreamWriter, @@ -123,6 +125,17 @@ describe(`RecordBatchStreamReader over message-sized chunks`, () => { }); }); +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) {