feat(arrow): stream a result off the socket - #142
Conversation
`get_result_arrow` and `stream_result_arrow` both collect the whole response body before decoding, so peak memory scales with the size of the result. Neither is usable for a result larger than memory: the CLI is OOM-killed around 10M rows fetching one. `open_result_arrow` returns an `ArrowResultStream` that pulls body chunks as batches are asked for and feeds them to arrow-ipc's push decoder, so peak memory is one record batch whatever the result size. A reqwest chunk converts to an arrow `Buffer` without copying. Additive: both existing entry points keep their current behaviour and their documented connection semantics. The trade for not buffering is that the pooled connection stays checked out until the stream is drained or dropped. The schema is resolved when the stream is opened, so a caller writing a header row does not have to read data to learn the column names.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Approved. Three non-blocking nits are inline.
The decode pump is correct. StreamDecoder::decode returns Ok(None) only after draining the buffer it is given, so replacing pending instead of appending never drops bytes. The split-at-every-offset test covers the chunk-boundary risk directly.
`next_batch` marked the stream done before confirming a clean end of stream, so a caller that logged the decode error and read on was told the stream had ended normally. A truncated download then reads as a complete result on the second call. `finish()` now runs first and `done` is set only once it succeeds. An exhausted body keeps yielding `None`, so every later call reaches `finish()` again and repeats the error. Also documents the new entry point where a reader would look for it: the module doc still said "two entry points are provided", the README's Arrow section described only the two buffering methods, and CHANGELOG had no entry under [Unreleased].
There was a problem hiding this comment.
All three prior threads are resolved in the code: finish() now runs before done is set, the module doc and README list the third entry point, and the CHANGELOG has an [Unreleased] entry. Two nits inline, neither blocking.
Note on CI: at review time the unit-test and integration checks were queued or running, so the test results are not visible here.
`ResultError::TooLarge` told the caller to stream the result with `Client::stream_result_arrow`. That method collects the whole body before decoding, so a caller who trips the auto-materialize guard on a huge result and follows the advice materializes exactly what the guard refused. It now names `open_result_arrow`, which decodes off the socket. Also covers `ArrowResultStream::read_all`, which nothing called: a dropped field there would hand back metadata-free results on every call.
| /// A truncated body must keep failing. If the stream marked itself done | ||
| /// before confirming a clean end, a caller that logged the first error and | ||
| /// read on would be told the stream ended normally — a short download | ||
| /// silently becoming a complete result. | ||
| /// `read_all` must carry every batch *and* the metadata headers into the | ||
| /// returned `ArrowResult`. A dropped field here would hand back | ||
| /// metadata-free results on every call. |
There was a problem hiding this comment.
super nit: split the merged doc comment (not blocking).
Two unrelated test docs now sit in one block on streaming_reader_read_all_carries_batches_and_metadata. The first three sentences describe a_cut_short_body_keeps_erroring_on_every_later_call at src/arrow.rs:1122, which is left undocumented. A reader of this block learns the wrong reason for the test below it.
| /// A truncated body must keep failing. If the stream marked itself done | |
| /// before confirming a clean end, a caller that logged the first error and | |
| /// read on would be told the stream ended normally — a short download | |
| /// silently becoming a complete result. | |
| /// `read_all` must carry every batch *and* the metadata headers into the | |
| /// returned `ArrowResult`. A dropped field here would hand back | |
| /// metadata-free results on every call. | |
| /// `read_all` must carry every batch *and* the metadata headers into the | |
| /// returned `ArrowResult`. A dropped field here would hand back | |
| /// metadata-free results on every call. |
Then move the truncated-body sentences onto a_cut_short_body_keeps_erroring_on_every_later_call.
There was a problem hiding this comment.
Cycle 3 review. Both prior nits are addressed: the TooLarge doc and message now name open_result_arrow with the reason stream_result_arrow is wrong there, and read_all has a test covering batches, schema, total_row_count and next_link. No blocking issues. One super nit inline on a merged test doc comment.
The integration job is still in progress, so this approval does not cover its result.
1 participant
Why
Both existing Arrow entry points buffer the entire response body before decoding:
get_result_arrow— "Buffers the entire Arrow IPC stream into memory before returning."stream_result_arrow— decodes lazily, but its own doc notes "the response body is collected once when the stream is created".So peak memory scales with the result, and neither is usable for a result larger
than memory. The hotdata CLI is OOM-killed (SIGKILL, no error message) around
10 M rows fetching one.
What
open_result_arrowreturns anArrowResultStreamthat pulls body chunks asbatches are asked for, feeding
arrow_ipc::reader::StreamDecoder. Peak memory isone record batch regardless of result size.
bytes::Bytes->arrow_buffer::Bufferis a refcount bump, not a copy.
The schema is resolved at open time, so a caller writing a CSV header does not
have to read data first.
Additive.
get_result_arrowandstream_result_arroware untouched,including their documented "connection released at construction" semantics. The
trade for not buffering is that
open_result_arrowholds the pooled connectionuntil the stream is drained or dropped.
Tests
20 arrow tests pass, 8 of them new. The ones worth noting:
streaming_reader_survives_a_split_at_every_offset— splits the IPC body atevery byte offset in turn and asserts identical batches. Chunk boundaries
landing inside a message are the whole risk of a push decoder, and a socket
only produces them occasionally.
streaming_reader_matches_the_buffered_reader— byte-for-byte equivalence withStreamReaderover the same bytes.a_body_cut_short_is_an_error_not_a_short_read— a truncated download must notread as a complete result. The buffered path got this for free; a streaming
reader has to assert it.
The byte source is swappable behind a
#[cfg(test)]enum so the pump can bedriven with chosen splits rather than whatever chunking a socket happens to
produce.
Merge order
This releases first — the CLI change that consumes it (hotdata-cli#TBD) is
pinned to a path override until there is a published version.