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	placement_for(queue, ordered_by_index.map(|index| stored_ordered_hash(shape, row, index)), row_number)
33}
34
35#[allow(clippy::disallowed_methods)]
36pub fn stored_ordered_hash(shape: &RowShape, row: &EncodedQueueRow, index: usize) -> Partition {
37	Partition::of(&[shape.get_value(row.as_slice(), index)])
38}
39
40pub fn placement_for(queue: &Queue, ordered: Option<Partition>, row_number: RowNumber) -> QueuePlacement {
41	#[allow(clippy::disallowed_methods)]
42	let hash = ordered.unwrap_or_else(|| Partition::of(&[Value::Uint8(row_number.0)]));
43	placement_from_hash(hash, queue.partitions(), ordered.is_some())
44}
45
46fn placement_from_hash(hash: Partition, partitions: u16, keyed: bool) -> QueuePlacement {
47	QueuePlacement {
48		partition: (hash.0 % partitions as u128) as u16,
49		key_hash: keyed.then_some(hash.0 as u64),
50	}
51}
52
53#[cfg(test)]
54mod tests {
55	use super::*;
56
57	#[test]
58	fn test_partition_is_the_full_hash_modulo_the_partition_count() {
59		let hash = Partition(0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210);
60
61		let placement = placement_from_hash(hash, 1000, true);
62
63		assert_eq!(placement.partition, 40);
64		assert_eq!(((hash.0 as u64) % 1000) as u16, 720, "the truncated hash would have placed it elsewhere");
65	}
66
67	#[test]
68	fn test_a_keyed_queue_truncates_the_hash_and_an_unkeyed_one_reports_no_key() {
69		let hash = Partition(0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210);
70
71		assert_eq!(placement_from_hash(hash, 16, true).key_hash, Some(0xFEDC_BA98_7654_3210));
72		assert_eq!(placement_from_hash(hash, 16, false).key_hash, None);
73	}
74
75	#[test]
76	fn test_a_single_partition_queue_places_everything_in_partition_zero() {
77		for hash in [0u128, 1, u128::MAX] {
78			assert_eq!(placement_from_hash(Partition(hash), 1, true).partition, 0);
79		}
80	}
81}