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

reifydb_engine/
partition.rs
