reifydb_core/key/
partition.rs1use reifydb_codec::key::{
5 deserializer::KeyDeserializer,
6 encoded::{EncodedKey, EncodedKeyRange},
7 serializer::KeySerializer,
8};
9use reifydb_value::value::partition::Partition;
10
11use super::{EncodableKey, KeyKind};
12use crate::{
13 interface::catalog::object::ObjectId,
14 key::catalog::{KeyDeserializerCatalogExt, KeySerializerCatalogExt},
15};
16
17#[derive(Debug, Clone, PartialEq)]
18pub struct PartitionKey {
19 pub object: ObjectId,
20 pub partition: Partition,
21}
22
23impl PartitionKey {
24 pub fn new(object: impl Into<ObjectId>, partition: Partition) -> Self {
25 Self {
26 object: object.into(),
27 partition,
28 }
29 }
30
31 pub fn encoded(object: impl Into<ObjectId>, partition: Partition) -> EncodedKey {
32 Self::new(object, partition).encode()
33 }
34
35 pub fn full_scan(object: impl Into<ObjectId>) -> EncodedKeyRange {
36 let object = object.into();
37 let mut start = KeySerializer::with_capacity(10);
38 start.extend_u8(Self::KIND as u8).extend_object_id(object);
39 let mut end = KeySerializer::with_capacity(10);
40 end.extend_u8(Self::KIND as u8).extend_object_id(object.prev());
41 EncodedKeyRange::start_end(Some(start.to_encoded_key()), Some(end.to_encoded_key()))
42 }
43}
44
45impl EncodableKey for PartitionKey {
46 const KIND: KeyKind = KeyKind::Partition;
47
48 fn encode(&self) -> EncodedKey {
49 let mut serializer = KeySerializer::with_capacity(26);
50 serializer.extend_u8(Self::KIND as u8).extend_object_id(self.object).extend_u128(self.partition.0);
51 serializer.to_encoded_key()
52 }
53
54 fn decode(key: &EncodedKey) -> Option<Self> {
55 let mut de = KeyDeserializer::from_bytes(key.as_slice());
56 let kind: KeyKind = de.read_u8().ok()?.try_into().ok()?;
57 if kind != Self::KIND {
58 return None;
59 }
60 let object = de.read_object_id().ok()?;
61 let partition = Partition(de.read_u128().ok()?);
62 Some(Self {
63 object,
64 partition,
65 })
66 }
67}
68
69#[cfg(test)]
70mod tests {
71 use std::ops::RangeBounds;
72
73 use reifydb_value::value::{Value, partition::Partition};
74
75 use super::{EncodableKey, PartitionKey};
76 use crate::interface::catalog::{id::TableId, object::ObjectId};
77
78 #[test]
79 fn test_roundtrip() {
80 let key = PartitionKey {
81 object: ObjectId::Table(TableId(7)),
82 partition: Partition::of(&[Value::Utf8("us".to_string())]),
83 };
84 let decoded = PartitionKey::decode(&key.encode()).unwrap();
85 assert_eq!(decoded, key);
86 }
87
88 #[test]
89 fn test_partitions_of_object_share_prefix() {
90 let object = ObjectId::Table(TableId(3));
91 let range = PartitionKey::full_scan(object);
92 let k = PartitionKey::encoded(object, Partition::of(&[Value::Utf8("us".to_string())]));
93 assert!(range.contains(&k));
94 let other = PartitionKey::encoded(
95 ObjectId::Table(TableId(4)),
96 Partition::of(&[Value::Utf8("us".to_string())]),
97 );
98 assert!(!range.contains(&other));
99 }
100}