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

reifydb_engine/
partition.rs
