Skip to main content

reifydb_engine/queue/
partition.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use reifydb_codec::row::{queue::EncodedQueueRow, shape::RowShape};
5use reifydb_core::{interface::catalog::queue::Queue, internal_error};
6use reifydb_value::value::{Value, partition::Partition, row_number::RowNumber};
7
8use crate::Result;
9
10pub struct QueuePlacement {
11	pub partition: u16,
12	pub key_hash: Option<u64>,
13}
14
15pub fn ordered_by_index(queue: &Queue) -> Result<Option<usize>> {
16	let Some(ordered_by) = queue.ordered_by() else {
17		return Ok(None);
18	};
19	let index = queue.columns.iter().position(|c| c.name == *ordered_by).ok_or_else(|| {
20		internal_error!("queue {} declares ordered_by {} which is not a column", queue.name, ordered_by)
21	})?;
22	Ok(Some(index))
23}
24
25pub fn placement_of(
26	queue: &Queue,
27	shape: &RowShape,
28	row: &EncodedQueueRow,
29	ordered_by_index: Option<usize>,
30	row_number: RowNumber,
31) -> QueuePlacement {
32	let hash = match ordered_by_index {
33		Some(index) => Partition::of(&[shape.get_value(row.as_slice(), index)]),
34		None => Partition::of(&[Value::Uint8(row_number.0)]),
35	};
36	placement_from_hash(hash, queue.partitions(), ordered_by_index.is_some())
37}
38
39fn placement_from_hash(hash: Partition, partitions: u16, keyed: bool) -> QueuePlacement {
40	QueuePlacement {
41		partition: (hash.0 % partitions as u128) as u16,
42		key_hash: keyed.then_some(hash.0 as u64),
43	}
44}
45
46#[cfg(test)]
47mod tests {
48	use super::*;
49
50	#[test]
51	fn test_partition_is_the_full_hash_modulo_the_partition_count() {
52		let hash = Partition(0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210);
53
54		let placement = placement_from_hash(hash, 1000, true);
55
56		assert_eq!(placement.partition, 40);
57		assert_eq!(((hash.0 as u64) % 1000) as u16, 720, "the truncated hash would have placed it elsewhere");
58	}
59
60	#[test]
61	fn test_a_keyed_queue_truncates_the_hash_and_an_unkeyed_one_reports_no_key() {
62		let hash = Partition(0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210);
63
64		assert_eq!(placement_from_hash(hash, 16, true).key_hash, Some(0xFEDC_BA98_7654_3210));
65		assert_eq!(placement_from_hash(hash, 16, false).key_hash, None);
66	}
67
68	#[test]
69	fn test_a_single_partition_queue_places_everything_in_partition_zero() {
70		for hash in [0u128, 1, u128::MAX] {
71			assert_eq!(placement_from_hash(Partition(hash), 1, true).partition, 0);
72		}
73	}
74}