reifydb_engine/
partition.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{collections::HashSet, sync::LazyLock};
5
6use postcard::to_stdvec;
7use reifydb_codec::row::shape::{RowFamily, RowShape, RowShapeField};
8use reifydb_core::{
9	interface::catalog::{id::TableId, object::ObjectId, table::Table},
10	key::{
11		any::TaggedKey,
12		partition::PartitionKey,
13		row::{PartitionedRowKey, RowKey},
14	},
15	partition::{PartitionError, partition_col_indices},
16};
17use reifydb_transaction::transaction::Transaction;
18use reifydb_value::value::{Value, blob::Blob, partition::Partition, row_number::RowNumber, value_type::ValueType};
19
20use crate::Result;
21
22static REGISTRY_SHAPE: LazyLock<RowShape> =
23	LazyLock::new(|| RowShape::new(RowFamily::Pod, vec![RowShapeField::unconstrained("values", ValueType::Blob)]));
24
25pub fn partition_values(shape: &RowShape, row: &[u8], indices: &[usize]) -> Vec<Value> {
26	indices.iter().map(|&i| shape.get_value(row, i)).collect()
27}
28
29pub fn table_partition_of_row(table: &Table, shape: &RowShape, row: &[u8]) -> Partition {
30	let indices = partition_col_indices(&table.columns, &table.partition_by);
31	Partition::of(&partition_values(shape, row, &indices))
32}
33
34pub fn table_row_key(table: &Table, shape: &RowShape, row: &[u8], row_number: RowNumber) -> TaggedKey {
35	if table.partition_by.is_empty() {
36		RowKey::new(table.id, row_number).into()
37	} else {
38		let partition = table_partition_of_row(table, shape, row);
39		PartitionedRowKey::new(table.id, partition, row_number).into()
40	}
41}
42
43pub fn row_key_from_partition(table_id: TableId, partition: Option<Partition>, row_number: RowNumber) -> TaggedKey {
44	match partition {
45		None => RowKey::new(table_id, row_number).into(),
46		Some(partition) => PartitionedRowKey::new(table_id, partition, row_number).into(),
47	}
48}
49
50pub fn resolve_partition(
51	txn: &mut Transaction<'_>,
52	object: ObjectId,
53	partition: Partition,
54	values: &[Value],
55	verified: &mut HashSet<Partition>,
56) -> Result<()> {
57	if !verified.insert(partition) {
58		return Ok(());
59	}
60	let key = PartitionKey::new(object, partition);
61	let encoded = to_stdvec(values).expect("value postcard is total");
62	let candidate = Value::Blob(Blob::from(encoded));
63	match txn.get(&key)? {
64		Some(multi) => {
65			if REGISTRY_SHAPE.get_value(&multi.bytes, 0) != candidate {
66				return Err(PartitionError::PartitionHashCollision {
67					object,
68					hash: partition.0,
69				}
70				.into());
71			}
72		}
73		None => {
74			let mut row = REGISTRY_SHAPE.allocate_pod();
75			REGISTRY_SHAPE.set_value(&mut row, 0, &candidate);
76			txn.set(&key, row.freeze())?;
77		}
78	}
79	Ok(())
80}
reifydb_engine/partition.rs

reifydb_engine/
partition.rs
