Skip to main content

laminar_core/shuffle/
message.rs

1//! Logical messages carried by the ordered shuffle transport.
2//!
3//! The data/barrier wire encoding is gRPC/protobuf (`proto/shuffle.proto`, `ShuffleFrame`);
4//! transport identity is internal to [`super::transport`]. A data message's batch is
5//! Arrow IPC-encoded as one self-contained logical payload. Wire fragmentation
6//! slices that allocation without copying it or retaining per-stage codec state.
7
8use std::sync::Arc;
9
10use arrow_array::RecordBatch;
11
12use crate::checkpoint::barrier::CheckpointBarrier;
13
14/// Maximum reassembled Arrow IPC payload for one logical shuffled batch: 16 MiB.
15pub const MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024;
16
17/// Logical message carried on a shuffle connection.
18#[derive(Debug, Clone, PartialEq)]
19pub enum ShuffleMessage {
20    /// An aligned checkpoint barrier ordered after preceding checkpointed data.
21    Barrier(CheckpointBarrier),
22    /// A stage batch with a non-empty canonical route set.
23    Data {
24        /// Stable stage demultiplexing scope.
25        stage: String,
26        /// Ascending, duplicate-free receiver-owned vnode route set.
27        routed_vnodes: Arc<[u32]>,
28        /// User batch, with its schema left unchanged.
29        batch: RecordBatch,
30    },
31}
32
33impl ShuffleMessage {
34    /// Construct stateful data covered by checkpoint delivery accounting.
35    #[must_use]
36    pub fn checkpointed(stage: String, vnode: u32, batch: RecordBatch) -> Self {
37        Self::checkpointed_routed(stage, Arc::from([vnode]), batch)
38    }
39
40    /// Construct owner-coalesced stateful data with canonical out-of-band vnode metadata.
41    #[must_use]
42    pub fn checkpointed_routed(
43        stage: String,
44        routed_vnodes: Arc<[u32]>,
45        batch: RecordBatch,
46    ) -> Self {
47        Self::Data {
48            stage,
49            routed_vnodes,
50            batch,
51        }
52    }
53}