reifydb_engine/
partition.rs1use std::collections::HashMap;
5
6use reifydb_core::{
7 interface::catalog::{id::TableId, object::ObjectId},
8 key::{
9 any::TaggedKey,
10 partition::PartitionKey,
11 row::{PartitionedRowKey, RowKey},
12 },
13 partition::{PartitionError, registry_matches, registry_row},
14};
15use reifydb_transaction::transaction::Transaction;
16use reifydb_value::{
17 error::Error,
18 value::{Value, partition::Partition, row_number::RowNumber},
19};
20
21use crate::Result;
22
23pub fn row_key_from_partition(table_id: TableId, partition: Option<Partition>, row_number: RowNumber) -> TaggedKey {
24 match partition {
25 None => RowKey::new(table_id, row_number).into(),
26 Some(partition) => PartitionedRowKey::new(table_id, partition, row_number).into(),
27 }
28}
29
30pub fn resolve_partitions(
31 txn: &mut Transaction<'_>,
32 object: ObjectId,
33 distinct: &[(Partition, Vec<Value>)],
34 verified: &mut HashMap<Partition, Vec<Value>>,
35) -> Result<()> {
36 for (partition, values) in distinct {
37 match verified.get(partition) {
38 Some(known) if known != values => return Err(collision(object, *partition)),
39 Some(_) => {}
40 None => {
41 check_or_register(txn, object, *partition, values)?;
42 verified.insert(*partition, values.clone());
43 }
44 }
45 }
46 Ok(())
47}
48
49fn check_or_register(
50 txn: &mut Transaction<'_>,
51 object: ObjectId,
52 partition: Partition,
53 values: &[Value],
54) -> Result<()> {
55 let key = PartitionKey::new(object, partition);
56 match txn.get(&[key.clone().into()])?.pop().flatten() {
57 Some(multi) if !registry_matches(&multi.bytes, values) => Err(collision(object, partition)),
58 Some(_) => Ok(()),
59 None => txn.set(vec![(key.clone().into(), registry_row(values))]),
60 }
61}
62
63fn collision(object: ObjectId, partition: Partition) -> Error {
64 PartitionError::PartitionHashCollision {
65 object,
66 hash: partition.0,
67 }
68 .into()
69}