use std::{
collections::{BTreeSet, HashMap},
sync::Arc,
};
use arrow_array::RecordBatch;
use reifydb_codec::{
key::encoded::EncodedKey,
row::{
bytes::EncodedBytes,
pod::EncodedPodRow,
shape::RowShape,
table::{EncodedTableRow, EncodedTableRowBuilder},
},
};
use reifydb_core::{
error::diagnostic::{
catalog::{namespace_not_found, table_not_found},
engine,
index::primary_key_violation,
query::column_not_found,
},
interface::{
catalog::{
config::{ConfigKey, GetConfig},
id::IndexId,
key::PrimaryKey,
namespace::Namespace,
object::ObjectId,
policy::{DataOp, PolicyTargetType},
table::Table,
},
resolved::{ResolvedNamespace, ResolvedObject, ResolvedTable},
},
key::{any::TaggedKey, catalog::IndexEntryKey},
partition::PartitionError,
value::batch::single_row,
};
use reifydb_evaluate::stack::SymbolTable;
use reifydb_policy::enforce::CompiledWritePolicies;
use reifydb_rql::nodes::UpdateTableNode;
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
fragment::Fragment,
params::Params,
return_error,
value::{
Value,
identity::IdentityId,
partition::Partition,
row_number::RowNumber,
system_columns::{partitions, row_numbers, user_columns},
},
};
use super::{
columns::{ColumnPipeline, input_views, intern_dictionary_columns},
context::{TableTarget, WriteExecCtx},
primary_key::{self, PrimaryKeyEncoder},
returning::{
decode_returning_dictionaries, decode_rows_to_columns, decode_stored_table_rows, evaluate_returning,
with_pre_image,
},
shape::get_or_create_table_shape,
table_insert::cast_partitions,
};
use crate::{
Result,
error::EngineError,
partition::row_key_from_partition,
policy::PolicyEvaluator,
transaction::operation::table::{TableOperations, load_stored_shape},
vm::{
instruction::dml::{
coerce::InputFragments,
time::{RowTime, populator_index},
},
services::Services,
volcano::{
compile::compile,
query::{QueryContext, QueryNode, query_budget},
},
},
};
pub(crate) fn update_table(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
plan: UpdateTableNode,
params: Params,
symbols: &SymbolTable,
) -> Result<RecordBatch> {
let UpdateTableNode {
input,
target,
returning,
} = plan;
let target = target.expect("Cannot infer target table from pipeline - no table found");
let (namespace, table) = resolve_update_table_target(services, txn, &target)?;
let shape = get_or_create_table_shape(&services.catalog, &table, txn)?;
let target_data = TableTarget {
namespace: &namespace,
table: &table,
fragment: target.identifier(),
};
let context = build_update_table_query_context(services, &target_data, ¶ms, symbols, txn.identity());
let fragments = InputFragments::of(&input);
let mut input_node = compile(*input, txn, Arc::new(context.clone()));
input_node.initialize(txn, &context)?;
let exec = WriteExecCtx {
services,
symbols,
};
let (updated_count, returned_rows, pre_rows) = run_table_update(
&exec,
txn,
&mut input_node,
&fragments,
&target_data,
&shape,
&context,
returning.is_some(),
)?;
if let Some(returning_exprs) = &returning {
let columns = decode_rows_to_columns(&shape, &returned_rows)?;
let columns = decode_returning_dictionaries(services, txn, &table.columns, columns)?;
let pre_columns = decode_stored_table_rows(&services.catalog, txn, &table, &shape, &pre_rows)?;
let pre_columns = decode_returning_dictionaries(services, txn, &table.columns, pre_columns)?;
let columns = with_pre_image(columns, &pre_columns)?;
return evaluate_returning(services, symbols, returning_exprs, columns, txn.identity());
}
update_table_result(namespace.name(), &table.name, updated_count)
}
#[inline]
fn resolve_update_table_target(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
target: &ResolvedTable,
) -> Result<(Namespace, Table)> {
let namespace_name = target.namespace().name();
let Some(namespace) = services.catalog.find_namespace_by_name(txn, namespace_name)? else {
return_error!(namespace_not_found(Fragment::internal(namespace_name), namespace_name));
};
let Some(table) = services.catalog.find_table_by_name(txn, namespace.id(), target.name())? else {
let fragment = target.identifier().clone();
return_error!(table_not_found(fragment.clone(), namespace_name, target.name(),));
};
Ok((namespace, table))
}
#[inline]
fn build_update_table_query_context(
services: &Arc<Services>,
target: &TableTarget<'_>,
params: &Params,
symbols: &SymbolTable,
identity: IdentityId,
) -> QueryContext {
let namespace_ident = Fragment::internal(target.namespace.name());
let resolved_namespace = ResolvedNamespace::new(namespace_ident, target.namespace.clone());
let table_ident = Fragment::internal(target.table.name.clone());
let resolved_table = ResolvedTable::new(table_ident, resolved_namespace, target.table.clone());
QueryContext {
services: services.clone(),
source: Some(ResolvedObject::Table(resolved_table)),
batch_size: services.catalog.get_config_uint2(ConfigKey::QueryRowBatchSize) as u64,
params: params.clone(),
symbols: symbols.clone(),
identity,
memory: query_budget(services),
}
}
type ReturnedRows = Vec<(RowNumber, EncodedBytes)>;
#[allow(clippy::too_many_arguments)]
fn run_table_update(
exec: &WriteExecCtx<'_>,
txn: &mut Transaction<'_>,
input_node: &mut Box<dyn QueryNode>,
fragments: &InputFragments,
target: &TableTarget<'_>,
shape: &RowShape,
context: &QueryContext,
has_returning: bool,
) -> Result<(u64, ReturnedRows, ReturnedRows)> {
let mut updated_count = 0u64;
let mut returned_rows: ReturnedRows = Vec::new();
let mut pre_rows: ReturnedRows = Vec::new();
let mut mutable_context = context.clone();
let pipeline = ColumnPipeline {
columns: &target.table.columns,
sequences: None,
series_key: None,
fragments,
context,
};
let pk_def = primary_key::get_primary_key(&exec.services.catalog, txn, target.table)?;
let mut encoder: Option<PrimaryKeyEncoder> = None;
let populator = populator_index(&target.table.time, shape);
let evaluator = PolicyEvaluator::new(exec.services, exec.symbols);
let mut policies: Option<Option<CompiledWritePolicies>> = None;
while let Some(columns) = input_node.next(txn, &mut mutable_context)? {
if columns.num_rows() == 0 {
continue;
}
if policies.is_none() {
policies = Some(evaluator.compile_write_policies(
txn,
target.namespace.name(),
&target.table.name,
DataOp::Update,
PolicyTargetType::Table,
)?);
}
let compiled = policies.as_ref().and_then(Option::as_ref);
if let Some(compiled) = compiled {
evaluator.enforce_compiled(compiled, txn.identity(), &columns)?;
}
if let Some((unknown, _)) = user_columns(&columns)
.find(|(field, _)| !target.table.columns.iter().any(|c| &c.name == field.name()))
{
return_error!(column_not_found(fragments.column(unknown.name())));
}
if row_numbers(&columns)?.is_empty() {
return_error!(engine::missing_row_number_column());
}
let partitioned = !target.table.partition_by.is_empty();
if partitioned && partitions(&columns)?.len() != columns.num_rows() {
return Err(EngineError::MissingPartitionAddress {
object: ObjectId::Table(target.table.id),
operation: "UPDATE",
}
.into());
}
let row_numbers: Vec<RowNumber> = row_numbers(&columns)?.to_vec();
let sidecar_partitions: Vec<Partition> = partitions(&columns)?;
let row_count = columns.num_rows();
let row_keys: Vec<TaggedKey> = row_numbers
.iter()
.enumerate()
.map(|(row_idx, &row_number)| {
row_key_from_partition(
target.table.id,
sidecar_partitions.get(row_idx).copied(),
row_number,
)
})
.collect();
let old_rows: Vec<EncodedBytes> = txn
.get(&row_keys)?
.into_iter()
.map(|row| row.expect("bytes must exist for update").bytes)
.collect();
enforce_old_row_policies(exec, txn, target, shape, &row_numbers, &old_rows, compiled)?;
let inputs = input_views(&columns, &target.table.columns)?;
let mut batches = [pipeline.cast_target_columns(&inputs, row_count, None)?];
intern_dictionary_columns(
&exec.services.catalog,
txn,
pipeline.columns,
pipeline.series_key,
&mut batches,
)?;
let mut rows: Vec<EncodedTableRowBuilder> = (0..row_count).map(|_| shape.allocate_table()).collect();
batches[0].write(shape, &mut rows)?;
let now = exec.services.runtime_context.clock.now();
let times = RowTime::new(
&target.table.name,
&target.table.time,
shape,
populator.map(|index| batches[0].view(index)).transpose()?,
);
let new_partitions =
cast_partitions(&target.table.columns, &target.table.partition_by, &batches[0])?.rows;
let mut rotation = KeyRotation::default();
let mut prepared_rows: Vec<EncodedTableRowBuilder> = Vec::with_capacity(row_count);
let mut partitions_out: Vec<Partition> = Vec::with_capacity(row_count);
let mut pre_by_row: HashMap<RowNumber, EncodedBytes> = HashMap::new();
for (row_idx, (mut row, &row_number)) in rows.into_iter().zip(row_numbers.iter()).enumerate() {
let partition = sidecar_partitions.get(row_idx).copied();
if let Some(old) = partition {
if new_partitions[row_idx] != old {
return Err(PartitionError::ImmutablePartitionColumn {
object: ObjectId::Table(target.table.id),
}
.into());
}
}
let old_row = &old_rows[row_idx];
if let Some(pk_def) = &pk_def {
let encoder = match &mut encoder {
Some(encoder) => encoder,
None => encoder.insert(PrimaryKeyEncoder::new(pk_def, target.table)?),
};
let fingerprint = EncodedTableRow::view(old_row).fingerprint();
let loaded;
let pre_shape = if fingerprint == shape.fingerprint() {
shape
} else {
loaded = load_stored_shape(
&exec.services.catalog,
txn,
target.table,
fingerprint,
)?;
&loaded
};
rotation.rotate(
txn, target, shape, pre_shape, pk_def, encoder, old_row, &row, row_number,
)?;
}
if has_returning {
pre_by_row.insert(row_number, old_row.clone());
}
let old_row = EncodedTableRow::view(old_row);
let old_created_at = old_row.created_at();
let old_time = old_row.time();
row.set_timestamps(old_created_at, now);
if let Some(time) = times.updated(row_idx, &row, old_time)? {
row.set_time(time);
}
prepared_rows.push(row);
if let Some(p) = partition {
partitions_out.push(p);
}
}
rotation.write(txn)?;
let key_columns = encoder.as_ref().map(PrimaryKeyEncoder::key_columns).unwrap_or_default();
let stored = txn.update_table(
target.table,
&row_numbers,
&partitions_out,
&mut prepared_rows,
&key_columns,
&old_rows,
&exec.services.catalog,
)?;
updated_count += stored.len() as u64;
if has_returning {
for (row_number, _) in &stored {
let pre = pre_by_row
.remove(row_number)
.expect("every stored row must carry the pre image read before its write");
pre_rows.push((*row_number, pre));
}
returned_rows.extend(stored);
}
}
Ok((updated_count, returned_rows, pre_rows))
}
fn enforce_old_row_policies(
exec: &WriteExecCtx<'_>,
txn: &mut Transaction<'_>,
target: &TableTarget<'_>,
shape: &RowShape,
row_numbers: &[RowNumber],
rows: &[EncodedBytes],
compiled: Option<&CompiledWritePolicies>,
) -> Result<()> {
let Some(compiled) = compiled.filter(|compiled| compiled.reads_row()) else {
return Ok(());
};
let old_rows: ReturnedRows = row_numbers.iter().copied().zip(rows.iter().cloned()).collect();
let old_columns = decode_stored_table_rows(&exec.services.catalog, txn, target.table, shape, &old_rows)?;
let old_columns = decode_returning_dictionaries(exec.services, txn, &target.table.columns, old_columns)?;
PolicyEvaluator::new(exec.services, exec.symbols).enforce_compiled_row_dependent(
compiled,
txn.identity(),
&old_columns,
)
}
#[derive(Default)]
struct KeyRotation {
removed: BTreeSet<EncodedKey>,
added: BTreeSet<EncodedKey>,
removes: Vec<(TaggedKey, Option<EncodedBytes>)>,
sets: Vec<(TaggedKey, EncodedBytes)>,
}
impl KeyRotation {
#[allow(clippy::too_many_arguments)]
fn rotate(
&mut self,
txn: &mut Transaction<'_>,
target: &TableTarget<'_>,
shape: &RowShape,
pre_shape: &RowShape,
pk_def: &PrimaryKey,
encoder: &PrimaryKeyEncoder,
pre_row: &[u8],
new_row: &[u8],
row_number: RowNumber,
) -> Result<()> {
let index = IndexId::primary(pk_def.id);
let pre_key: TaggedKey =
IndexEntryKey::new(target.table.id, index, encoder.encode(pre_shape, pre_row)).into();
let post_key: TaggedKey =
IndexEntryKey::new(target.table.id, index, encoder.encode(shape, new_row)).into();
let pre_encoded = pre_key.encode();
let post_encoded = post_key.encode();
if pre_encoded != post_encoded {
self.added.remove(&pre_encoded);
self.removed.insert(pre_encoded);
self.removes.push((pre_key, None));
let taken = if self.added.contains(&post_encoded) {
true
} else if self.removed.contains(&post_encoded) {
false
} else {
txn.contains(&[post_key.clone()])?[0]
};
if taken {
let key_columns = pk_def.columns.iter().map(|c| c.name.clone()).collect();
return_error!(primary_key_violation(
target.fragment.clone(),
target.table.name.clone(),
key_columns
));
}
}
self.removed.remove(&post_encoded);
self.added.insert(post_encoded);
self.sets.push((post_key, EncodedPodRow::new(&u64::from(row_number).to_be_bytes()).into_bytes()));
Ok(())
}
fn write(self, txn: &mut Transaction<'_>) -> Result<()> {
txn.remove(self.removes)?;
txn.set(self.sets)
}
}
#[inline]
fn update_table_result(namespace: &str, table: &str, updated: u64) -> Result<RecordBatch> {
single_row([
("namespace", Value::Utf8(namespace.to_string())),
("table", Value::Utf8(table.to_string())),
("updated", Value::Uint8(updated)),
])
}