reifydb_core/key/
partition.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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::shape::ShapeId,
14	key::catalog::{KeyDeserializerCatalogExt, KeySerializerCatalogExt},
15};
16
17#[derive(Debug, Clone, PartialEq)]
18pub struct PartitionKey {
19	pub shape: ShapeId,
20	pub partition: Partition,
21}
22
23impl PartitionKey {
24	pub fn new(shape: impl Into<ShapeId>, partition: Partition) -> Self {
25		Self {
26			shape: shape.into(),
27			partition,
28		}
29	}
30
31	pub fn encoded(shape: impl Into<ShapeId>, partition: Partition) -> EncodedKey {
32		Self::new(shape, partition).encode()
33	}
34
35	pub fn full_scan(shape: impl Into<ShapeId>) -> EncodedKeyRange {
36		let shape = shape.into();
37		let mut start = KeySerializer::with_capacity(10);
38		start.extend_u8(Self::KIND as u8).extend_shape_id(shape);
39		let mut end = KeySerializer::with_capacity(10);
40		end.extend_u8(Self::KIND as u8).extend_shape_id(shape.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_shape_id(self.shape).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 shape = de.read_shape_id().ok()?;
61		let partition = Partition(de.read_u128().ok()?);
62		Some(Self {
63			shape,
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, shape::ShapeId};
77
78	#[test]
79	fn test_roundtrip() {
80		let key = PartitionKey {
81			shape: ShapeId::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_shape_share_prefix() {
90		let shape = ShapeId::Table(TableId(3));
91		let range = PartitionKey::full_scan(shape);
92		let k = PartitionKey::encoded(shape, Partition::of(&[Value::Utf8("us".to_string())]));
93		assert!(range.contains(&k));
94		let other = PartitionKey::encoded(
95			ShapeId::Table(TableId(4)),
96			Partition::of(&[Value::Utf8("us".to_string())]),
97		);
98		assert!(!range.contains(&other));
99	}
100}
reifydb_core/key/partition.rs

reifydb_core/key/
partition.rs
