laminar_core/shuffle/routing/
mod.rs

1//! Row-to-vnode routing shared by checkpointed cluster shuffle paths.
2
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use arrow::compute::take;
7use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
8
9use crate::state::partition_key::{PartitionKeyCodecV1Builder, MAX_PARTITION_KEY_COLUMNS};
10use crate::state::{NodeId, PartitionKeyCodecError, PartitionKeyCodecV1, VnodeAssignmentSnapshot};
11
12/// Target decoded size for one routed batch. This is intentionally below the hard receiver bound
13/// so schema/allocator variance does not turn ordinary skew into a transport failure.
14pub const ROUTE_TARGET_BATCH_BYTES: usize = 4 * 1024 * 1024;
15/// Hard decoded size admitted for one logical shuffle batch.
16pub const ROUTE_MAX_BATCH_BYTES: usize = 8 * 1024 * 1024;
17/// Row-count bound prevents tiny or zero-width rows from producing unbounded logical frames.
18pub const ROUTE_MAX_BATCH_ROWS: usize = 65_536;
19
20/// Logical Arrow bytes referenced by this slice, independent of backing-buffer capacity or owner.
21/// This is the stable bound for IPC content; transport reservations account the retained backing
22/// allocation separately.
23///
24/// # Errors
25///
26/// Returns an Arrow compute error when buffer-size accounting overflows.
27pub fn logical_batch_bytes(batch: &RecordBatch) -> Result<usize, arrow_schema::ArrowError> {
28    batch.columns().iter().try_fold(0usize, |total, column| {
29        let data = column.to_data();
30        let bytes = data
31            .get_slice_memory_size()?
32            .checked_add(variadic_buffer_bytes(&data)?)
33            .ok_or_else(|| {
34                arrow_schema::ArrowError::ComputeError(
35                    "logical shuffle batch size overflow".to_string(),
36                )
37            })?;
38        total.checked_add(bytes).ok_or_else(|| {
39            arrow_schema::ArrowError::ComputeError(
40                "logical shuffle batch size overflow".to_string(),
41            )
42        })
43    })
44}
45
46fn variadic_buffer_bytes(
47    data: &arrow::array::ArrayData,
48) -> Result<usize, arrow_schema::ArrowError> {
49    let own = if matches!(
50        data.data_type(),
51        arrow_schema::DataType::Utf8View | arrow_schema::DataType::BinaryView
52    ) {
53        data.buffers()
54            .iter()
55            .skip(1)
56            .try_fold(0usize, |total, buffer| {
57                total.checked_add(buffer.len()).ok_or_else(|| {
58                    arrow_schema::ArrowError::ComputeError(
59                        "variadic Arrow buffer size overflow".to_string(),
60                    )
61                })
62            })?
63    } else {
64        0
65    };
66    data.child_data().iter().try_fold(own, |total, child| {
67        total
68            .checked_add(variadic_buffer_bytes(child)?)
69            .ok_or_else(|| {
70                arrow_schema::ArrowError::ComputeError(
71                    "nested variadic Arrow buffer size overflow".to_string(),
72                )
73            })
74    })
75}
76
77/// A local vnode slice in deterministic vnode order.
78#[derive(Debug, Clone)]
79pub struct LocalRoute {
80    /// Exact local vnode.
81    pub vnode: u32,
82    /// Rows for that vnode, preserving their input order.
83    pub batch: RecordBatch,
84    /// Original input-row indices, aligned one-for-one with `batch`.
85    pub source_rows: Arc<[u32]>,
86}
87
88/// One bounded owner-coalesced remote batch.
89#[derive(Debug, Clone)]
90pub struct RemoteRoute {
91    /// Certified remote owner.
92    pub owner: NodeId,
93    /// Ascending, duplicate-free vnode set represented by `batch`.
94    pub routed_vnodes: Arc<[u32]>,
95    /// Rows for this owner, preserving their input order.
96    pub batch: RecordBatch,
97}
98
99/// Complete, conserving routing plan staged before any remote send.
100#[derive(Debug, Clone, Default)]
101pub struct CheckpointRoutePlan {
102    /// Local chunks, ordered by vnode and then input row.
103    pub local: Vec<LocalRoute>,
104    /// Remote chunks, ordered by owner and then input row.
105    pub remote: Vec<RemoteRoute>,
106}
107
108/// A routing failure. No plan is returned, so callers cannot partially send or silently drop rows.
109#[derive(Debug, thiserror::Error)]
110pub enum ShuffleRoutingError {
111    /// Keyed routing requires at least one key column.
112    #[error("shuffle routing requires at least one key column")]
113    EmptyKey,
114    /// Modulo routing is undefined for an empty vnode space.
115    #[error("shuffle routing requires a nonzero vnode count")]
116    EmptyVnodeSpace,
117    /// A resolved key index no longer exists in the batch schema.
118    #[error("shuffle key column {index} is outside the {columns}-column batch")]
119    KeyColumnOutOfRange {
120        /// Requested zero-based index.
121        index: usize,
122        /// Available columns.
123        columns: usize,
124    },
125    /// ABI v1 admits only scalar, non-floating key types.
126    #[error("shuffle key column {index} has unsupported partition type {data_type}")]
127    UnsupportedKeyType {
128        /// Requested zero-based index.
129        index: usize,
130        /// Rejected Arrow type.
131        data_type: arrow_schema::DataType,
132    },
133    /// A bounded partition-key ABI invariant failed before row encoding.
134    #[error("shuffle partition-key contract: {0}")]
135    PartitionKeyContract(PartitionKeyCodecError),
136    /// Arrow row encoding or slicing failed.
137    #[error("shuffle Arrow routing: {0}")]
138    Arrow(#[from] arrow_schema::ArrowError),
139    /// Routing metadata must cover every input row exactly once.
140    #[error("shuffle route has {vnodes} vnode entries for {rows} rows")]
141    RowVnodeCardinality {
142        /// Input row count.
143        rows: usize,
144        /// Supplied route count.
145        vnodes: usize,
146    },
147    /// `UInt32` take indices cannot represent this input cardinality.
148    #[error("shuffle batch has {rows} rows; maximum supported is {}", u32::MAX)]
149    TooManyRows {
150        /// Input row count.
151        rows: usize,
152    },
153    /// The route references a vnode outside the pinned assignment.
154    #[error("shuffle vnode {vnode} is outside the {vnode_count}-vnode assignment")]
155    VnodeOutOfRange {
156        /// Invalid vnode.
157        vnode: u32,
158        /// Pinned assignment cardinality.
159        vnode_count: usize,
160    },
161    /// Formation/rebalance has not assigned this vnode yet.
162    #[error("shuffle vnode {vnode} is unassigned")]
163    UnassignedVnode {
164        /// Vnode awaiting an owner.
165        vnode: u32,
166    },
167    /// A single row cannot be split further and would always fail receiver admission.
168    #[error("shuffle row for vnode {vnode} occupies {bytes} decoded bytes; hard limit is {limit}")]
169    OversizedRow {
170        /// Row's routed vnode.
171        vnode: u32,
172        /// Decoded Arrow bytes.
173        bytes: usize,
174        /// Hard decoded limit.
175        limit: usize,
176    },
177}
178
179impl ShuffleRoutingError {
180    /// Whether retrying after assignment convergence can succeed without changing the input.
181    #[must_use]
182    pub const fn is_not_ready(&self) -> bool {
183        matches!(self, Self::UnassignedVnode { .. })
184    }
185}
186
187/// Hash each row with the engine's canonical Arrow-row and xxh3 encoding.
188///
189/// # Errors
190/// Returns a structural or Arrow encoding error instead of panicking on invalid resolved columns.
191pub fn row_vnodes(
192    batch: &RecordBatch,
193    columns: &[usize],
194    vnode_count: u32,
195) -> Result<Vec<u32>, ShuffleRoutingError> {
196    if columns.is_empty() {
197        return Err(ShuffleRoutingError::EmptyKey);
198    }
199    if vnode_count == 0 {
200        return Err(ShuffleRoutingError::EmptyVnodeSpace);
201    }
202    if columns.len() > MAX_PARTITION_KEY_COLUMNS {
203        return Err(ShuffleRoutingError::PartitionKeyContract(
204            PartitionKeyCodecError::TooManyKeyColumns {
205                count: columns.len(),
206                limit: MAX_PARTITION_KEY_COLUMNS,
207            },
208        ));
209    }
210    let mut cols: Vec<ArrayRef> = Vec::with_capacity(columns.len());
211    let mut builder = PartitionKeyCodecV1Builder::with_capacity(columns.len());
212    for &index in columns {
213        let column = batch.columns().get(index).cloned().ok_or(
214            ShuffleRoutingError::KeyColumnOutOfRange {
215                index,
216                columns: batch.num_columns(),
217            },
218        )?;
219        builder
220            .push(column.data_type().clone())
221            .map_err(|error| match error {
222                PartitionKeyCodecError::EmptyKeySchema => ShuffleRoutingError::EmptyKey,
223                PartitionKeyCodecError::UnsupportedKeyType { data_type, .. } => {
224                    ShuffleRoutingError::UnsupportedKeyType { index, data_type }
225                }
226                PartitionKeyCodecError::Arrow(error) => ShuffleRoutingError::Arrow(error),
227                other => ShuffleRoutingError::PartitionKeyContract(other),
228            })?;
229        cols.push(column);
230    }
231    let codec = builder.finish().map_err(|error| match error {
232        PartitionKeyCodecError::EmptyKeySchema => ShuffleRoutingError::EmptyKey,
233        PartitionKeyCodecError::UnsupportedKeyType { index, data_type } => {
234            ShuffleRoutingError::UnsupportedKeyType {
235                index: columns[index],
236                data_type,
237            }
238        }
239        PartitionKeyCodecError::Arrow(error) => ShuffleRoutingError::Arrow(error),
240        other => ShuffleRoutingError::PartitionKeyContract(other),
241    })?;
242    let rows = codec.encode_columns(&cols)?;
243    let vnode_count =
244        std::num::NonZeroU32::new(vnode_count).ok_or(ShuffleRoutingError::EmptyVnodeSpace)?;
245    Ok((0..batch.num_rows())
246        .map(|row| PartitionKeyCodecV1::vnode_for_encoded(rows.row(row).as_ref(), vnode_count))
247        .collect())
248}
249
250/// Build a complete local/remote plan from one caller-pinned assignment.
251///
252/// Every input row is assigned exactly once or the whole call fails. Remote batches carry vnode
253/// ownership out-of-band; user schemas, including a user field named `__laminar_vnode`, are never
254/// rewritten.
255///
256/// # Errors
257/// Returns before producing a plan for malformed metadata, unassigned/out-of-range vnodes, Arrow
258/// slicing failures, or a single row above the receiver's hard decoded-memory bound.
259pub fn route_checkpointed_batch(
260    batch: &RecordBatch,
261    row_vnodes: &[u32],
262    assignment: &VnodeAssignmentSnapshot,
263    self_id: NodeId,
264) -> Result<CheckpointRoutePlan, ShuffleRoutingError> {
265    if row_vnodes.len() != batch.num_rows() {
266        return Err(ShuffleRoutingError::RowVnodeCardinality {
267            rows: batch.num_rows(),
268            vnodes: row_vnodes.len(),
269        });
270    }
271    if batch.num_rows() > usize::try_from(u32::MAX).unwrap_or(usize::MAX) {
272        return Err(ShuffleRoutingError::TooManyRows {
273            rows: batch.num_rows(),
274        });
275    }
276    if batch.num_rows() == 0 {
277        return Ok(CheckpointRoutePlan::default());
278    }
279
280    let mut local_groups: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
281    let mut remote_groups: BTreeMap<NodeId, Vec<(u32, u32)>> = BTreeMap::new();
282    for (row, &vnode) in row_vnodes.iter().enumerate() {
283        let owner = assignment
284            .owners()
285            .get(usize::try_from(vnode).unwrap_or(usize::MAX))
286            .copied()
287            .ok_or(ShuffleRoutingError::VnodeOutOfRange {
288                vnode,
289                vnode_count: assignment.owners().len(),
290            })?;
291        if owner.is_unassigned() {
292            return Err(ShuffleRoutingError::UnassignedVnode { vnode });
293        }
294        let row = u32::try_from(row).map_err(|_| ShuffleRoutingError::TooManyRows {
295            rows: batch.num_rows(),
296        })?;
297        if owner == self_id {
298            local_groups.entry(vnode).or_default().push(row);
299        } else {
300            remote_groups.entry(owner).or_default().push((row, vnode));
301        }
302    }
303
304    let mut plan = CheckpointRoutePlan::default();
305    for (vnode, indices) in local_groups {
306        for (range, slice) in bounded_slices(batch, &indices, vnode)? {
307            plan.local.push(LocalRoute {
308                vnode,
309                batch: slice,
310                source_rows: Arc::from(&indices[range]),
311            });
312        }
313    }
314    for (owner, rows) in remote_groups {
315        let indices: Vec<u32> = rows.iter().map(|(row, _)| *row).collect();
316        for (range, slice) in bounded_slices(batch, &indices, rows[0].1)? {
317            let mut routed_vnodes: Vec<u32> = rows[range].iter().map(|(_, vnode)| *vnode).collect();
318            routed_vnodes.sort_unstable();
319            routed_vnodes.dedup();
320            plan.remote.push(RemoteRoute {
321                owner,
322                routed_vnodes: routed_vnodes.into(),
323                batch: slice,
324            });
325        }
326    }
327    debug_assert_eq!(
328        plan.local
329            .iter()
330            .map(|route| route.batch.num_rows())
331            .sum::<usize>()
332            + plan
333                .remote
334                .iter()
335                .map(|route| route.batch.num_rows())
336                .sum::<usize>(),
337        batch.num_rows()
338    );
339    Ok(plan)
340}
341
342fn bounded_slices(
343    batch: &RecordBatch,
344    indices: &[u32],
345    vnode: u32,
346) -> Result<Vec<(std::ops::Range<usize>, RecordBatch)>, ShuffleRoutingError> {
347    let estimated_row_bytes = logical_batch_bytes(batch)?
348        .div_ceil(batch.num_rows())
349        .max(1);
350    let initial_rows =
351        ROUTE_MAX_BATCH_ROWS.min((ROUTE_TARGET_BATCH_BYTES / estimated_row_bytes).max(1));
352    let mut chunks = Vec::new();
353    let mut start = 0;
354    while start < indices.len() {
355        let mut end = (start + initial_rows).min(indices.len());
356        loop {
357            let slice = take_rows(batch, &indices[start..end])?;
358            let bytes = logical_batch_bytes(&slice)?;
359            let rows = end - start;
360            if bytes <= ROUTE_TARGET_BATCH_BYTES || rows == 1 {
361                if bytes > ROUTE_MAX_BATCH_BYTES {
362                    return Err(ShuffleRoutingError::OversizedRow {
363                        vnode,
364                        bytes,
365                        limit: ROUTE_MAX_BATCH_BYTES,
366                    });
367                }
368                chunks.push((start..end, slice));
369                start = end;
370                break;
371            }
372            end = start + (rows / 2).max(1);
373        }
374    }
375    Ok(chunks)
376}
377
378fn take_rows(batch: &RecordBatch, indices: &[u32]) -> Result<RecordBatch, ShuffleRoutingError> {
379    if indices.len() == batch.num_rows() && batch.get_array_memory_size() <= ROUTE_MAX_BATCH_BYTES {
380        debug_assert!(indices
381            .iter()
382            .enumerate()
383            .all(|(offset, &index)| u32::try_from(offset).ok() == Some(index)));
384        return Ok(batch.clone());
385    }
386
387    let indices = UInt32Array::from(indices.to_vec());
388    let columns = batch
389        .columns()
390        .iter()
391        .map(|column| take(column, &indices, None))
392        .collect::<Result<Vec<_>, _>>()?;
393    Ok(RecordBatch::try_new(batch.schema(), columns)?)
394}
395
396#[cfg(test)]
397mod tests;
laminar_core/shuffle/routing/mod.rs

laminar_core/shuffle/routing/
mod.rs
