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 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}