alamb commented on code in PR #24592:
URL: https://github.com/apache/datafusion/pull/24592#discussion_r3846691267
##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -84,6 +91,213 @@ struct SpillReaderStream {
// Small margin allowed to accommodate slight memory accounting variation
const SPILL_BATCH_MEMORY_MARGIN: usize = 4096;
+/// Reassembles an IPC stream read in arbitrary chunks into one exactly sized
+/// allocation per message, so that the zero-copy [`StreamDecoder`] produces
+/// batches whose buffers pin only their own message.
+///
+/// # Why
+///
+/// The decoder builds arrays on slices of whatever [`Buffer`] it is given,
+/// and a slice keeps its whole backing allocation alive. Fed with the raw
+/// chunks of the byte stream (128 KB for a file backed spill) that goes
+/// wrong in two ways.
+///
+/// A message that fits inside a chunk pins the entire chunk. With ~5 KB
+/// batches, one chunk holds ~27 of them, and each decoded batch retains, and
+/// is accounted for, 128 KB:
+///
+/// ```text
+/// chunk (128 KB allocation)
+/// ┌──────┬──────┬──────┬─────┬───────┐
+/// │ msg1 │ msg2 │ msg3 │ ... │ msg27 │
+/// └──────┴──────┴──────┴─────┴───────┘
+/// ▲
+/// batch1's buffers slice here, yet keep all 128 KB alive
+/// ```
+///
+/// A message that spans two chunks cannot be sliced, so the decoder gathers
+/// it into a `Vec` grown by doubling, and the batch keeps the spare
+/// capacity: a 256 KB body typically lands in a 512 KB allocation.
+///
+/// ```text
+/// chunk N chunk N+1
+/// ┌──────┬────────────────────┬──────────────┬────────┐
+/// │ ... │ msgK (first part) │ msgK (rest) │ msgK+1 │
+/// └──────┴────────────────────┴──────────────┴────────┘
+/// ```
+///
+/// Either way a batch uses several times the memory recorded for it at
+/// spill time, breaking the `max_record_batch_memory` budgeting that the
+/// multi-level merge relies on.
+///
+/// # How
+///
+/// Each message is copied out of the chunks into allocations sized from its
+/// own headers: a head buffer (the length prefix and flatbuffer metadata,
+/// whose `bodyLength` gives the body size) and, when non-empty, a body
+/// buffer of exactly that size. The decoder then zero-copies from the body
+/// buffer, so a batch pins exactly its own message:
+///
+/// ```text
+/// body for msg1 (5 KB) body for msgK (256 KB)
+/// ┌──────┐ ┌────────────────────┐
+/// │ msg1 │ ◀── batch1 │ msgK │ ◀── batchK
+/// └──────┘ └────────────────────┘
+/// ```
+///
+/// This costs one copy per message, which the decoder already paid for
+/// spanning messages, without the doubling reallocation.
+struct MessageFramer {
+ state: FramerState,
+}
+
+enum FramerState {
+ /// Reading the 4 byte continuation marker or metadata length.
+ Prefix {
+ head: Vec<u8>,
+ read: usize,
+ continuation: bool,
+ },
+ /// Reading the flatbuffer metadata into `head`, which already holds the
+ /// prefix and is allocated for `head_len` bytes.
+ Metadata {
+ head: Vec<u8>,
+ metadata_start: usize,
+ head_len: usize,
+ },
+ /// Reading the body into `body`, which is allocated for `body_len` bytes.
+ Body {
+ head: Vec<u8>,
+ body: Vec<u8>,
+ body_len: usize,
+ },
+ /// The end-of-stream marker was read.
+ Finished,
+}
+
+impl MessageFramer {
+ fn new() -> Self {
+ Self {
+ state: FramerState::prefix(),
+ }
+ }
+
+ /// Consumes bytes from `input` until a message is complete or `input` is
+ /// exhausted, returning the buffers of a completed message.
+ fn push(&mut self, input: &mut Buffer) -> Result<Option<Vec<Buffer>>> {
+ while !input.is_empty() {
+ match &mut self.state {
+ FramerState::Prefix {
+ head,
+ read,
+ continuation,
+ } => {
+ let to_read = input.len().min(4 - *read);
+ head.extend_from_slice(&input[..to_read]);
+ input.advance(to_read);
+ *read += to_read;
+ if *read < 4 {
+ continue;
+ }
+ let word: [u8; 4] = head[head.len() -
4..].try_into().unwrap();
+ if !*continuation && word == CONTINUATION_MARKER {
+ *continuation = true;
+ *read = 0;
+ continue;
+ }
+ let metadata_len = u32::from_le_bytes(word) as usize;
+ let head = std::mem::take(head);
+ if metadata_len == 0 {
+ self.state = FramerState::Finished;
+ return Ok(Some(vec![Buffer::from_vec(head)]));
+ }
+ let metadata_start = head.len();
+ let head_len = metadata_start + metadata_len;
+ let mut sized = Vec::with_capacity(head_len);
+ sized.extend_from_slice(&head);
+ self.state = FramerState::Metadata {
+ head: sized,
+ metadata_start,
+ head_len,
+ };
+ }
+ FramerState::Metadata {
+ head,
+ metadata_start,
+ head_len,
+ } => {
+ let to_read = input.len().min(*head_len - head.len());
+ head.extend_from_slice(&input[..to_read]);
+ input.advance(to_read);
+ if head.len() < *head_len {
+ continue;
+ }
+ let message =
+ root_as_message(&head[*metadata_start..]).map_err(|e| {
+ datafusion_common::exec_datafusion_err!(
+ "Invalid IPC message in spill file: {e}"
+ )
+ })?;
+ let body_len =
+ usize::try_from(message.bodyLength()).map_err(|_| {
+ datafusion_common::exec_datafusion_err!(
+ "Invalid IPC message body length in spill
file: {}",
+ message.bodyLength()
+ )
+ })?;
+ let head = std::mem::take(head);
+ if body_len == 0 {
+ self.state = FramerState::prefix();
+ return Ok(Some(vec![Buffer::from_vec(head)]));
+ }
+ self.state = FramerState::Body {
+ head,
+ body: Vec::with_capacity(body_len),
+ body_len,
+ };
+ }
+ FramerState::Body {
+ head,
+ body,
+ body_len,
+ } => {
+ let to_read = input.len().min(*body_len - body.len());
+ body.extend_from_slice(&input[..to_read]);
Review Comment:
I think this line effectively copies each input byte twice (once into
`input` and then once into `body`)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
