reifydb_engine/transaction/operation/

table.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use arrow_array::RecordBatch;
use reifydb_catalog::catalog::Catalog;
use reifydb_codec::row::{
	bytes::{EncodedBytes, RowBuilder},
	shape::{RowFamily, RowShape, fingerprint::RowShapeFingerprint},
	table::{EncodedTableRow, EncodedTableRowBuilder},
};
use reifydb_core::{
	common::{ChangeVersion, CommitVersion},
	delta::SetPre,
	error::diagnostic::index::primary_key_changed_by_hook,
	interface::{
		catalog::{object::ObjectId, table::Table},
		change::{Change, ChangeOrigin, Diff},
	},
	internal_error,
	key::any::TaggedKey,
	partition::{PartitionError, partition_col_indices, partitions_of_rows},
	row::row_shape_from_columns,
	value::batch::{append_rows_fitted, from_encoded_bytes},
};
use reifydb_transaction::{
	change::{RowChange, TableRowInsertion},
	interceptor::{WithInterceptors, table_row::TableRowInterceptor},
	transaction::{Transaction, admin::AdminTransaction, command::CommandTransaction, read::Read, write::Write},
};
use reifydb_value::{
	fragment::Fragment,
	return_error,
	value::{Value, datetime::DateTime, partition::Partition, row_number::RowNumber},
};
use smallvec::smallvec;

use crate::{Result, partition::row_key_from_partition};

fn build_table_insert_change(
	table: &Table,
	shape: &RowShape,
	ids: &[RowNumber],
	bytes_slice: &[EncodedBytes],
) -> Result<Change> {
	Ok(Change {
		origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
		version: ChangeVersion::from(CommitVersion(0)),
		diffs: smallvec![Diff::insert(from_encoded_bytes(shape, ids, bytes_slice)?)],
		changed_at: DateTime::default(),
	})
}

fn build_table_update_change(
	table: &Table,
	shape: &RowShape,
	ids: &[RowNumber],
	pres: &[EncodedBytes],
	posts: &[EncodedBytes],
	stored_shape: impl FnMut(RowShapeFingerprint) -> Result<RowShape>,
) -> Result<Change> {
	Ok(Change {
		origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
		version: ChangeVersion::from(CommitVersion(0)),
		diffs: smallvec![Diff::update(
			decode_table_rows(shape, ids, pres, stored_shape)?,
			from_encoded_bytes(shape, ids, posts)?,
		)],
		changed_at: DateTime::default(),
	})
}

fn build_table_remove_change(
	table: &Table,
	shape: &RowShape,
	ids: &[RowNumber],
	bytes_slice: &[EncodedBytes],
	stored_shape: impl FnMut(RowShapeFingerprint) -> Result<RowShape>,
) -> Result<Change> {
	Ok(Change {
		origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
		version: ChangeVersion::from(CommitVersion(0)),
		diffs: smallvec![Diff::remove(decode_table_rows(shape, ids, bytes_slice, stored_shape)?)],
		changed_at: DateTime::default(),
	})
}

pub(crate) fn decode_table_rows(
	shape: &RowShape,
	ids: &[RowNumber],
	bytes_slice: &[EncodedBytes],
	mut stored_shape: impl FnMut(RowShapeFingerprint) -> Result<RowShape>,
) -> Result<RecordBatch> {
	let current = shape.fingerprint();
	if bytes_slice.iter().all(|row| EncodedTableRow::view(row).fingerprint() == current) {
		return from_encoded_bytes(shape, ids, bytes_slice);
	}
	let mut decoded = from_encoded_bytes(shape, &[], &[])?;
	let mut start = 0;
	while start < bytes_slice.len() {
		let fingerprint = EncodedTableRow::view(&bytes_slice[start]).fingerprint();
		let end = bytes_slice[start..]
			.iter()
			.position(|row| EncodedTableRow::view(row).fingerprint() != fingerprint)
			.map_or(bytes_slice.len(), |offset| start + offset);
		decoded = append_rows_fitted(
			decoded,
			&stored_shape(fingerprint)?,
			bytes_slice[start..end].iter().cloned(),
			ids[start..end].to_vec(),
		)?;
		start = end;
	}
	Ok(decoded)
}

pub(crate) fn load_stored_shape(
	catalog: &Catalog,
	txn: &mut Transaction<'_>,
	table: &Table,
	fingerprint: RowShapeFingerprint,
) -> Result<RowShape> {
	catalog.get_or_load_row_shape(fingerprint, txn)?.ok_or_else(|| {
		internal_error!("Row shape with fingerprint {:?} not found for table {}", fingerprint, table.name)
	})
}

pub struct TableInsert<'a> {
	pub ids: &'a [RowNumber],
	pub partitions: &'a [Partition],
	pub rows: &'a mut [EncodedTableRowBuilder],
	pub index: Vec<(TaggedKey, EncodedBytes)>,
	pub key_columns: &'a [usize],
}

fn insert_rows<T: WithInterceptors + Write>(
	txn: &mut T,
	table: &Table,
	shape: &RowShape,
	insert: TableInsert<'_>,
) -> Result<()> {
	let TableInsert {
		ids,
		partitions,
		rows,
		index,
		key_columns,
	} = insert;
	assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
	assert!(
		if table.partition_by.is_empty() {
			partitions.is_empty()
		} else {
			partitions.len() == ids.len()
		},
		"a partitioned table needs one partition per row, an unpartitioned table none"
	);
	if ids.is_empty() {
		return Ok(());
	}

	let hooked = !txn.table_row_pre_insert_interceptors().is_empty();
	let keys = if hooked {
		let keys = key_cells(shape, rows, key_columns);
		TableRowInterceptor::pre_insert(txn, table, ids, rows)?;
		Some(keys)
	} else {
		None
	};

	let frozen: Vec<EncodedBytes> = rows.iter().map(|row| row.clone().freeze_bytes()).collect();

	if let Some(keys) = keys {
		if moved_partitions(table, shape, partitions, &frozen)? {
			return Err(PartitionError::PartitionColumnChangedByHook {
				object: ObjectId::Table(table.id),
			}
			.into());
		}
		refuse_changed_keys(table, shape, rows, key_columns, &keys)?;
	}

	let mut entries: Vec<(TaggedKey, EncodedBytes, SetPre)> = Vec::with_capacity(frozen.len() + index.len());
	for (row_index, (row, &row_number)) in frozen.iter().zip(ids).enumerate() {
		entries.push((
			row_key_from_partition(table.id, partitions.get(row_index).copied(), row_number),
			row.clone(),
			SetPre::Absent,
		));
	}
	entries.extend(index.into_iter().map(|(key, bytes)| (key, bytes, SetPre::Unknown)));
	txn.set_with_pre(entries)?;

	TableRowInterceptor::post_insert(txn, table, ids, &frozen)?;

	let row_changes: Vec<RowChange> = ids
		.iter()
		.zip(frozen.iter())
		.map(|(&row_number, row)| {
			RowChange::TableInsert(TableRowInsertion {
				table_id: table.id,
				row_number,
				encoded: row.clone(),
			})
		})
		.collect();
	txn.track_row_change(&row_changes);

	let object = ObjectId::Table(table.id);
	txn.touch_flow_object(object);
	if txn.wants_flow_diff(object) {
		txn.track_flow_change(build_table_insert_change(table, shape, ids, &frozen)?);
	}

	Ok(())
}

fn key_cells(shape: &RowShape, rows: &[EncodedTableRowBuilder], key_columns: &[usize]) -> Vec<Value> {
	rows.iter().flat_map(|row| key_columns.iter().map(|&index| shape.get_value(row, index))).collect()
}

fn refuse_changed_keys(
	table: &Table,
	shape: &RowShape,
	rows: &[EncodedTableRowBuilder],
	key_columns: &[usize],
	keys: &[Value],
) -> Result<()> {
	if key_cells(shape, rows, key_columns) != keys {
		let key_columns = key_columns.iter().map(|&index| table.columns[index].name.clone()).collect();
		return_error!(primary_key_changed_by_hook(Fragment::None, table.name.clone(), key_columns));
	}
	Ok(())
}

fn moved_partitions(table: &Table, shape: &RowShape, partitions: &[Partition], rows: &[EncodedBytes]) -> Result<bool> {
	let indices = partition_col_indices(&table.columns, &table.partition_by);
	Ok(!indices.is_empty() && partitions_of_rows(shape, &indices, rows)? != partitions)
}

fn updated_rows<T: WithInterceptors>(
	txn: &mut T,
	table: &Table,
	shape: &RowShape,
	ids: &[RowNumber],
	partitions: &[Partition],
	rows: &mut [EncodedTableRowBuilder],
	key_columns: &[usize],
) -> Result<Vec<EncodedBytes>> {
	if txn.table_row_pre_update_interceptors().is_empty() {
		return Ok(rows.iter().map(|row| row.clone().freeze_bytes()).collect());
	}
	let keys = key_cells(shape, rows, key_columns);
	TableRowInterceptor::pre_update(txn, table, ids, rows)?;
	let posts: Vec<EncodedBytes> = rows.iter().map(|row| row.clone().freeze_bytes()).collect();
	if moved_partitions(table, shape, partitions, &posts)? {
		return Err(PartitionError::ImmutablePartitionColumn {
			object: ObjectId::Table(table.id),
		}
		.into());
	}
	refuse_changed_keys(table, shape, rows, key_columns, &keys)?;
	Ok(posts)
}

pub trait TableOperations {
	fn insert_table(&mut self, table: &Table, shape: &RowShape, insert: TableInsert<'_>) -> Result<()>;

	fn update_table(
		&mut self,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		rows: &mut [EncodedTableRowBuilder],
		key_columns: &[usize],
		pres: &[EncodedBytes],
		catalog: &Catalog,
	) -> Result<Vec<(RowNumber, EncodedBytes)>>;

	fn remove_from_table(
		&mut self,
		catalog: &Catalog,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		displayed: &[EncodedBytes],
	) -> Result<Vec<(RowNumber, EncodedBytes)>>;
}

fn update_rows<T: WithInterceptors + Write + Read>(
	txn: &mut T,
	table: &Table,
	ids: &[RowNumber],
	partitions: &[Partition],
	rows: &mut [EncodedTableRowBuilder],
	key_columns: &[usize],
	pres: &[EncodedBytes],
	mut stored_shape: impl FnMut(&mut T, RowShapeFingerprint) -> Result<RowShape>,
) -> Result<Vec<(RowNumber, EncodedBytes)>> {
	assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
	assert_eq!(ids.len(), pres.len(), "ids/pres length mismatch");
	if ids.is_empty() {
		return Ok(Vec::new());
	}

	let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
	let posts = updated_rows(txn, table, &shape, ids, partitions, rows, key_columns)?;

	let keys: Vec<TaggedKey> = ids
		.iter()
		.enumerate()
		.map(|(idx, &row_number)| row_key_from_partition(table.id, partitions.get(idx).copied(), row_number))
		.collect();
	let known: Vec<SetPre> = txn
		.get_committed(&keys)?
		.into_iter()
		.map(|committed| committed.map_or(SetPre::Unknown, |row| SetPre::Present(row.bytes)))
		.collect();
	txn.set_with_pre(
		keys.into_iter()
			.zip(posts.iter().cloned())
			.zip(known)
			.map(|((key, post), pre)| (key, post, pre))
			.collect(),
	)?;

	TableRowInterceptor::post_update(txn, table, ids, &posts, pres)?;

	let object = ObjectId::Table(table.id);
	txn.touch_flow_object(object);
	if txn.wants_flow_diff(object) {
		let change = build_table_update_change(table, &shape, ids, pres, &posts, |fingerprint| {
			stored_shape(txn, fingerprint)
		})?;
		txn.track_flow_change(change);
	}

	Ok(ids.iter().copied().zip(posts).collect())
}

fn remove_rows<T: WithInterceptors + Write + Read>(
	txn: &mut T,
	table: &Table,
	ids: &[RowNumber],
	partitions: &[Partition],
	displayed: &[EncodedBytes],
	mut stored_shape: impl FnMut(&mut T, RowShapeFingerprint) -> Result<RowShape>,
) -> Result<Vec<(RowNumber, EncodedBytes)>> {
	assert_eq!(ids.len(), displayed.len(), "ids/displayed length mismatch");
	if ids.is_empty() {
		return Ok(Vec::new());
	}

	let keys: Vec<TaggedKey> = ids
		.iter()
		.enumerate()
		.map(|(idx, &row_number)| row_key_from_partition(table.id, partitions.get(idx).copied(), row_number))
		.collect();
	let committed: Vec<Option<EncodedBytes>> =
		txn.get_committed(&keys)?.into_iter().map(|row| row.map(|v| v.bytes)).collect();
	let pre_for_cdc_bytes_vec: Vec<EncodedBytes> = committed
		.iter()
		.zip(displayed)
		.map(|(committed, displayed)| committed.clone().unwrap_or_else(|| displayed.clone()))
		.collect();

	TableRowInterceptor::pre_delete(txn, table, ids)?;

	txn.remove(keys.into_iter().zip(pre_for_cdc_bytes_vec.iter().cloned().map(Some)).collect())?;

	TableRowInterceptor::post_delete(txn, table, ids, &pre_for_cdc_bytes_vec)?;

	let object = ObjectId::Table(table.id);
	txn.touch_flow_object(object);
	if txn.wants_flow_diff(object) {
		let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
		let change = build_table_remove_change(table, &shape, ids, displayed, |fingerprint| {
			stored_shape(txn, fingerprint)
		})?;
		txn.track_flow_change(change);
	}

	Ok(ids.iter().copied().zip(displayed.iter().cloned()).collect())
}

impl TableOperations for CommandTransaction {
	fn insert_table(&mut self, table: &Table, shape: &RowShape, insert: TableInsert<'_>) -> Result<()> {
		insert_rows(self, table, shape, insert)
	}

	fn update_table(
		&mut self,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		rows: &mut [EncodedTableRowBuilder],
		key_columns: &[usize],
		pres: &[EncodedBytes],
		catalog: &Catalog,
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		update_rows(self, table, ids, partitions, rows, key_columns, pres, |txn, fingerprint| {
			load_stored_shape(catalog, &mut Transaction::Command(txn), table, fingerprint)
		})
	}

	fn remove_from_table(
		&mut self,
		catalog: &Catalog,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		displayed: &[EncodedBytes],
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		remove_rows(self, table, ids, partitions, displayed, |txn, fingerprint| {
			load_stored_shape(catalog, &mut Transaction::Command(txn), table, fingerprint)
		})
	}
}

impl TableOperations for AdminTransaction {
	fn insert_table(&mut self, table: &Table, shape: &RowShape, insert: TableInsert<'_>) -> Result<()> {
		insert_rows(self, table, shape, insert)
	}

	fn update_table(
		&mut self,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		rows: &mut [EncodedTableRowBuilder],
		key_columns: &[usize],
		pres: &[EncodedBytes],
		catalog: &Catalog,
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		update_rows(self, table, ids, partitions, rows, key_columns, pres, |txn, fingerprint| {
			load_stored_shape(catalog, &mut Transaction::Admin(txn), table, fingerprint)
		})
	}

	fn remove_from_table(
		&mut self,
		catalog: &Catalog,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		displayed: &[EncodedBytes],
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		remove_rows(self, table, ids, partitions, displayed, |txn, fingerprint| {
			load_stored_shape(catalog, &mut Transaction::Admin(txn), table, fingerprint)
		})
	}
}

impl TableOperations for Transaction<'_> {
	fn insert_table(&mut self, table: &Table, shape: &RowShape, insert: TableInsert<'_>) -> Result<()> {
		match self {
			Transaction::Command(txn) => txn.insert_table(table, shape, insert),
			Transaction::Admin(txn) => txn.insert_table(table, shape, insert),
			Transaction::Test(t) => t.inner.insert_table(table, shape, insert),
			Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
		}
	}

	fn update_table(
		&mut self,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		rows: &mut [EncodedTableRowBuilder],
		key_columns: &[usize],
		pres: &[EncodedBytes],
		catalog: &Catalog,
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		match self {
			Transaction::Command(txn) => {
				txn.update_table(table, ids, partitions, rows, key_columns, pres, catalog)
			}
			Transaction::Admin(txn) => {
				txn.update_table(table, ids, partitions, rows, key_columns, pres, catalog)
			}
			Transaction::Test(t) => {
				t.inner.update_table(table, ids, partitions, rows, key_columns, pres, catalog)
			}
			Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
		}
	}

	fn remove_from_table(
		&mut self,
		catalog: &Catalog,
		table: &Table,
		ids: &[RowNumber],
		partitions: &[Partition],
		displayed: &[EncodedBytes],
	) -> Result<Vec<(RowNumber, EncodedBytes)>> {
		match self {
			Transaction::Command(txn) => txn.remove_from_table(catalog, table, ids, partitions, displayed),
			Transaction::Admin(txn) => txn.remove_from_table(catalog, table, ids, partitions, displayed),
			Transaction::Test(t) => t.inner.remove_from_table(catalog, table, ids, partitions, displayed),
			Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
		}
	}
}