reifydb_engine/queue/
partition.rs1use 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}