use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions};
use arrow::datatypes::Schema;
use futures::StreamExt;
use tokio::task::JoinHandle;
use uuid::Uuid;
use crate::catalog::{
CatalogSchema, DataFileRecord, INTERNAL_ROW_ID_FIELD_NAME, InlineIndexRecord, RowValidity,
TransactionHelper, rows_to_record_batch,
};
use crate::expr::{
Expr, merge_filters, row_ids_in_list_expr, split_conjunction_filters, visited_columns,
};
use crate::storage::{Storage, read_data_file_by_record, read_row_id_array_from_data_file};
use crate::table::{IndexManager, Table, TableSchemaRef, process_insert_into_inline_rows_with_tx};
use crate::utils::{
extract_row_ids_from_record_batch, fixed_size_binary_array_to_uuids, timestamp_ms_from_now,
};
use crate::{ILError, ILResult, RecordBatchStream};
use std::time::Duration;
#[derive(Debug)]
pub struct TableUpdate {
pub set_map: HashMap<String, Expr>,
pub condition: Expr,
}
impl TableUpdate {
pub(crate) fn rewrite_columns(
mut self,
field_name_id_map: &HashMap<String, Uuid>,
) -> ILResult<Self> {
let mut new_set_map = HashMap::with_capacity(self.set_map.len());
for (key, value) in self.set_map {
let Some(field_id) = field_name_id_map.get(&key) else {
return Err(ILError::invalid_input(format!("Not found field {key}")));
};
let new_value = value.rewrite_columns(field_name_id_map)?;
new_set_map.insert(hex::encode(field_id), new_value);
}
let new_condition = self.condition.rewrite_columns(field_name_id_map)?;
self.set_map = new_set_map;
self.condition = new_condition;
Ok(self)
}
}
pub(crate) async fn process_update_by_condition(
tx_helper: &mut TransactionHelper,
table: &Table,
update: TableUpdate,
mut matched_data_file_rows: HashMap<Uuid, RecordBatchStream>,
) -> ILResult<usize> {
let updated_row_ids = update_inline_rows(tx_helper, table, &update).await?;
let inline_update_count = updated_row_ids.len();
if !updated_row_ids.is_empty() {
let updated_field_ids: Vec<String> = update.set_map.keys().cloned().collect();
let affected_index_ids = table
.index_manager
.index_ids_by_field_ids(&updated_field_ids);
if !affected_index_ids.is_empty() {
let updated_row_ids_set: HashSet<Uuid> = updated_row_ids.into_iter().collect();
for index_id in &affected_index_ids {
tx_helper
.invalidate_inline_index_rows(index_id, &updated_row_ids_set)
.await?;
}
let inline_index_records = build_inline_indexes_for_updated_rows(
tx_helper,
&table.table_id,
&table.table_schema,
&table.index_manager,
&updated_row_ids_set,
&affected_index_ids,
)
.await?;
tx_helper
.insert_inline_indexes(&inline_index_records)
.await?;
}
}
let data_file_records = tx_helper.get_data_files(&table.table_id).await?;
let mut file_update_count = 0;
for data_file_record in data_file_records {
let update_count =
if let Some(stream) = matched_data_file_rows.get_mut(&data_file_record.data_file_id) {
update_data_file_rows_by_matched_rows(
tx_helper,
table,
&update.set_map,
stream,
data_file_record,
)
.await?
} else {
update_data_file_rows_by_condition(
tx_helper,
table,
&update.set_map,
&update.condition,
data_file_record,
)
.await?
};
file_update_count += update_count;
}
Ok(inline_update_count + file_update_count)
}
pub(crate) async fn update_inline_rows(
tx_helper: &mut TransactionHelper,
table: &Table,
update: &TableUpdate,
) -> ILResult<Vec<Uuid>> {
let mut projection_columns = HashSet::new();
projection_columns.insert(INTERNAL_ROW_ID_FIELD_NAME.to_string());
projection_columns.extend(visited_columns(&update.condition));
for expr in update.set_map.values() {
projection_columns.extend(visited_columns(expr));
}
for col_name in update.set_map.keys() {
projection_columns.insert(col_name.clone());
}
let projected_fields: Vec<arrow::datatypes::Field> = table
.table_schema
.arrow_schema
.fields()
.iter()
.filter(|f| projection_columns.contains(f.name()))
.map(|f| (**f).clone())
.collect();
let projected_arrow_schema = Arc::new(Schema::new(projected_fields));
let projected_catalog_schema = Arc::new(CatalogSchema::from_arrow(&projected_arrow_schema)?);
let filters = split_conjunction_filters(vec![update.condition.clone()]);
let mut db_filters = Vec::new();
let mut arrow_filters = Vec::new();
for filter in filters {
if tx_helper
.catalog
.supports_filter(&filter, &table.table_schema.arrow_schema)?
{
db_filters.push(filter);
} else {
arrow_filters.push(filter);
}
}
let arrow_filter = merge_filters(arrow_filters);
let row_stream = tx_helper
.scan_inline_rows(
&table.table_id,
&projected_catalog_schema,
&db_filters,
None,
None,
)
.await?;
let mut chunk_stream = row_stream.chunks(100);
let mut matched_batches = Vec::new();
while let Some(row_chunk) = chunk_stream.next().await {
let rows = row_chunk.into_iter().collect::<ILResult<Vec<_>>>()?;
let record_batch = rows_to_record_batch(&projected_arrow_schema, &rows)?;
if let Some(arrow_filter) = &arrow_filter {
let bool_array = arrow_filter.condition_eval(&record_batch)?;
if bool_array.true_count() > 0 {
let filtered_batch =
arrow::compute::filter_record_batch(&record_batch, &bool_array)?;
matched_batches.push(filtered_batch);
}
} else {
matched_batches.push(record_batch);
}
}
drop(chunk_stream);
if matched_batches.is_empty() {
return Ok(Vec::new());
}
let mut updated_batches = Vec::new();
for batch in &matched_batches {
updated_batches.push(update_record_batch(batch, &update.set_map)?);
}
let mut all_row_ids = Vec::new();
let mut column_arrays: HashMap<&str, Vec<ArrayRef>> = HashMap::new();
for batch in &updated_batches {
let row_id_idx = batch.schema().index_of(INTERNAL_ROW_ID_FIELD_NAME)?;
let row_id_array = batch.column(row_id_idx);
let fixed_array = row_id_array
.as_any()
.downcast_ref::<arrow::array::FixedSizeBinaryArray>()
.ok_or_else(|| ILError::internal("Expected FixedSizeBinary row_id column"))?;
for i in 0..fixed_array.len() {
all_row_ids.push(Uuid::from_slice(fixed_array.value(i))?);
}
for col_name in update.set_map.keys() {
let idx = batch.schema().index_of(col_name)?;
column_arrays
.entry(col_name.as_str())
.or_default()
.push(batch.column(idx).clone());
}
}
let column_names: Vec<String> = update.set_map.keys().cloned().collect();
let mut concatenated_arrays = Vec::new();
for col_name in &column_names {
let arrays = column_arrays
.get(col_name.as_str())
.ok_or_else(|| ILError::internal(format!("Missing updated column: {col_name}")))?;
let array_refs: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
let concatenated = arrow::compute::concat(&array_refs)?;
concatenated_arrays.push(concatenated);
}
tx_helper
.update_inline_rows_by_row_ids(
&table.table_id,
&all_row_ids,
&column_names,
&concatenated_arrays,
)
.await?;
Ok(all_row_ids)
}
pub(crate) async fn update_data_file_rows_by_matched_rows(
tx_helper: &mut TransactionHelper,
table: &Table,
set_map: &HashMap<String, Expr>,
matched_data_file_rows: &mut RecordBatchStream,
data_file_record: DataFileRecord,
) -> ILResult<usize> {
let mut updated_row_ids = HashSet::new();
while let Some(batch) = matched_data_file_rows.next().await {
let batch = batch?;
if batch.num_rows() == 0 {
continue;
}
let row_ids = extract_row_ids_from_record_batch(&batch)?;
updated_row_ids.extend(row_ids);
let updated_batch = update_record_batch(&batch, set_map)?;
process_insert_into_inline_rows_with_tx(tx_helper, table, &[updated_batch]).await?;
}
let row_id_array = read_row_id_array_from_data_file(
table.storage.as_ref(),
&data_file_record.relative_path,
data_file_record.format,
)
.await?;
let row_ids = fixed_size_binary_array_to_uuids(&row_id_array)?;
tx_helper
.update_data_file_rows_as_invalid(
&data_file_record.data_file_id,
data_file_record.validity,
&row_ids,
&updated_row_ids,
)
.await?;
Ok(updated_row_ids.len())
}
pub(crate) async fn update_data_file_rows_by_condition(
tx_helper: &mut TransactionHelper,
table: &Table,
set_map: &HashMap<String, Expr>,
condition: &Expr,
data_file_record: DataFileRecord,
) -> ILResult<usize> {
let mut stream = read_data_file_by_record(
table.storage.as_ref(),
&table.table_schema,
&data_file_record,
None,
vec![condition.clone()],
1024,
)
.await?;
let mut updated_row_ids = HashSet::new();
while let Some(batch) = stream.next().await {
let batch = batch?;
let row_ids = extract_row_ids_from_record_batch(&batch)?;
updated_row_ids.extend(row_ids);
let updated_batch = update_record_batch(&batch, set_map)?;
process_insert_into_inline_rows_with_tx(tx_helper, table, &[updated_batch]).await?;
}
let row_id_array = read_row_id_array_from_data_file(
table.storage.as_ref(),
&data_file_record.relative_path,
data_file_record.format,
)
.await?;
let row_ids = fixed_size_binary_array_to_uuids(&row_id_array)?;
tx_helper
.update_data_file_rows_as_invalid(
&data_file_record.data_file_id,
data_file_record.validity,
&row_ids,
&updated_row_ids,
)
.await?;
Ok(updated_row_ids.len())
}
pub(crate) async fn parallel_find_matched_data_file_rows(
storage: Arc<dyn Storage>,
table_schema: TableSchemaRef,
condition: Expr,
data_file_records: Vec<DataFileRecord>,
) -> ILResult<HashMap<Uuid, RecordBatchStream>> {
let mut handles = Vec::new();
for data_file_record in data_file_records {
let storage = storage.clone();
let table_schema = table_schema.clone();
let condition = condition.clone();
let handle: JoinHandle<ILResult<(Uuid, RecordBatchStream)>> = tokio::spawn(async move {
let mut stream = read_data_file_by_record(
storage.as_ref(),
&table_schema,
&data_file_record,
None,
vec![condition],
1024,
)
.await?;
let mut prefetch_row_count = 0;
let mut stream_exhausted = true;
let mut batches = Vec::new();
while let Some(batch) = stream.next().await {
let batch = batch?;
prefetch_row_count += batch.num_rows();
batches.push(batch);
if prefetch_row_count > 1000 {
stream_exhausted = false;
break;
}
}
let memory_stream =
Box::pin(futures::stream::iter(batches).map(Ok)) as RecordBatchStream;
let merged_stream = if stream_exhausted {
memory_stream
} else {
Box::pin(futures::stream::select_all(vec![memory_stream, stream]))
};
Ok((data_file_record.data_file_id, merged_stream))
});
handles.push(handle);
}
let mut matched_rows = HashMap::new();
for handle in handles {
let (data_file_id, stream) = handle.await??;
matched_rows.insert(data_file_id, stream);
}
Ok(matched_rows)
}
fn update_record_batch(
batch: &RecordBatch,
set_map: &HashMap<String, Expr>,
) -> ILResult<RecordBatch> {
let mut columns = batch.columns().to_vec();
for (name, value) in set_map {
let idx = batch.schema().index_of(name)?;
let new_array = value.eval(batch)?.into_array(batch.num_rows())?;
columns[idx] = Arc::new(new_array);
}
let options = RecordBatchOptions::default().with_row_count(Some(batch.num_rows()));
Ok(RecordBatch::try_new_with_options(
batch.schema(),
columns,
&options,
)?)
}
async fn build_inline_indexes_for_updated_rows(
tx_helper: &mut TransactionHelper,
table_id: &Uuid,
table_schema: &TableSchemaRef,
index_manager: &IndexManager,
updated_row_ids: &HashSet<Uuid>,
affected_index_ids: &[Uuid],
) -> ILResult<Vec<InlineIndexRecord>> {
let catalog_schema = Arc::new(CatalogSchema::from_arrow(&table_schema.arrow_schema)?);
let row_ids_vec: Vec<Uuid> = updated_row_ids.iter().copied().collect();
let row_id_filter = row_ids_in_list_expr(row_ids_vec);
let row_stream = tx_helper
.scan_inline_rows(
table_id,
&catalog_schema,
std::slice::from_ref(&row_id_filter),
None,
None,
)
.await?;
let mut all_rows = Vec::new();
let mut chunk_stream = row_stream.chunks(1000);
while let Some(row_chunk) = chunk_stream.next().await {
let rows = row_chunk.into_iter().collect::<ILResult<Vec<_>>>()?;
all_rows.extend(rows);
}
if all_rows.is_empty() {
return Ok(Vec::new());
}
let record_batch = rows_to_record_batch(&table_schema.arrow_schema, &all_rows)?;
let mut inline_index_records = Vec::new();
for (index_def, index_kind) in index_manager.iter_index_and_kind() {
if !affected_index_ids.contains(&index_def.index_id) {
continue;
}
let mut index_builder = index_kind.builder(index_def)?;
index_builder.append(&record_batch)?;
let mut index_data = Vec::new();
index_builder.write_bytes(&mut index_data)?;
let row_ids: Vec<Uuid> = all_rows
.iter()
.filter_map(|row| row.get_row_id().ok().flatten())
.collect();
inline_index_records.push(InlineIndexRecord {
index_id: index_def.index_id,
created_at: timestamp_ms_from_now(Duration::ZERO),
row_ids,
validity: RowValidity::new(all_rows.len()),
index_data,
});
}
Ok(inline_index_records)
}