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 data and frontiers.
21    Barrier(CheckpointBarrier),
22    /// Per-stage event-time progress ordered with the stage's data.
23    Frontier {
24        /// Stable stage demultiplexing scope.
25        stage: String,
26        /// Current watermark, or `None` while the channel is uninitialized.
27        watermark: Option<i64>,
28        /// Whether the channel is excluded from downstream watermark minima.
29        idle: bool,
30    },
31    /// A stage batch with a non-empty canonical route set.
32    Data {
33        /// Stable stage demultiplexing scope.
34        stage: String,
35        /// Ascending, duplicate-free receiver-owned vnode route set.
36        routed_vnodes: Arc<[u32]>,
37        /// User batch, with its schema left unchanged.
38        batch: RecordBatch,
39    },
40}
41
42impl ShuffleMessage {
43    /// Construct stateful data covered by checkpoint delivery accounting.
44    #[must_use]
45    pub fn checkpointed(stage: String, vnode: u32, batch: RecordBatch) -> Self {
46        Self::checkpointed_routed(stage, Arc::from([vnode]), batch)
47    }
48
49    /// Construct owner-coalesced stateful data with canonical out-of-band vnode metadata.
50    #[must_use]
51    pub fn checkpointed_routed(
52        stage: String,
53        routed_vnodes: Arc<[u32]>,
54        batch: RecordBatch,
55    ) -> Self {
56        Self::Data {
57            stage,
58            routed_vnodes,
59            batch,
60        }
61    }
62}