use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use arrow::datatypes::SchemaRef;
use arrow_schema::Schema;
use futures::StreamExt;
use uuid::Uuid;
use crate::utils::timestamp_ms_from_now;
use crate::catalog::{
CatalogSchema, FieldRecord, IndexFileRecord, IndexRecord, InlineIndexRecord, RowValidity,
TableRecord, TransactionHelper, rows_to_record_batch,
};
use crate::expr::Expr;
use crate::index::{IndexConfig, IndexDefinition, IndexDefinitionRef, IndexKind, IndexParams};
use crate::storage::read_data_file_by_record;
use crate::table::{Table, TableConfig, TableSchemaRef, check_default_expr};
use crate::{ILError, ILResult, check_schema_contains_system_column};
#[derive(Debug, Clone)]
pub struct TableCreation {
pub namespace_name: String,
pub table_name: String,
pub schema: SchemaRef,
pub default_values: HashMap<String, Expr>,
pub config: TableConfig,
pub if_not_exists: bool,
}
impl TableCreation {
pub fn validate(&self) -> ILResult<()> {
if self.table_name.is_empty() {
return Err(ILError::invalid_input(
"Table name cannot be empty".to_string(),
));
}
if self.schema.fields().is_empty() {
return Err(ILError::invalid_input(
"Table schema cannot be empty".to_string(),
));
}
let field_names = self
.schema
.fields()
.iter()
.map(|f| f.name())
.collect::<HashSet<_>>();
if field_names.len() < self.schema.fields.len() {
return Err(ILError::invalid_input(
"Table schema contains duplicate field names".to_string(),
));
}
check_schema_contains_system_column(&self.schema)?;
Ok(())
}
}
impl Default for TableCreation {
fn default() -> Self {
Self {
namespace_name: String::default(),
table_name: String::default(),
schema: Arc::new(Schema::empty()),
default_values: HashMap::default(),
config: TableConfig::default(),
if_not_exists: false,
}
}
}
pub(crate) async fn process_create_table(
tx_helper: &mut TransactionHelper,
creation: TableCreation,
) -> ILResult<Uuid> {
let namespace_id = tx_helper
.get_namespace_id(&creation.namespace_name)
.await?
.ok_or_else(|| {
ILError::invalid_input(format!("Namespace {} not found", creation.namespace_name))
})?;
if let Some(table_id) = tx_helper
.get_table_id(&namespace_id, &creation.table_name)
.await?
{
return if creation.if_not_exists {
Ok(table_id)
} else {
Err(ILError::invalid_input(format!(
"Table {} already exists in namespace {}",
creation.table_name, creation.namespace_name
)))
};
}
for (field_name, default_value) in &creation.default_values {
let field = creation.schema.field_with_name(field_name)?;
check_default_expr(field, default_value, creation.schema.as_ref())?;
}
let table_id = Uuid::now_v7();
tx_helper
.insert_table(&TableRecord {
table_id,
table_name: creation.table_name,
namespace_id,
config: creation.config,
schema_metadata: creation.schema.metadata().clone(),
})
.await?;
let mut field_records = Vec::new();
for field in creation.schema.fields() {
let default_value = creation.default_values.get(field.name()).cloned();
field_records.push(FieldRecord::new(
Uuid::now_v7(),
table_id,
field,
default_value,
));
}
tx_helper.insert_fields(&field_records).await?;
tx_helper
.create_inline_row_table(&table_id, &field_records)
.await?;
Ok(table_id)
}
#[derive(Debug, Clone)]
pub struct IndexCreation {
pub name: String,
pub kind: String,
pub key_columns: Vec<String>,
pub params: Arc<dyn IndexParams>,
pub concurrency: usize,
pub if_not_exists: bool,
pub config: IndexConfig,
}
impl IndexCreation {
pub(crate) fn rewrite_columns(
mut self,
field_name_id_map: &HashMap<String, Uuid>,
) -> ILResult<Self> {
let mut rewritten_columns = Vec::with_capacity(self.key_columns.len());
for key_col in self.key_columns {
if let Some(field_id) = field_name_id_map.get(&key_col) {
rewritten_columns.push(hex::encode(field_id));
} else {
return Err(ILError::invalid_input(format!(
"key column {key_col} not found"
)));
}
}
self.key_columns = rewritten_columns;
Ok(self)
}
}
pub(crate) async fn process_create_index(
tx_helper: &mut TransactionHelper,
table: Table,
creation: IndexCreation,
) -> ILResult<Uuid> {
let index_id = Uuid::now_v7();
let index_kind = table
.index_manager
.get_index_kind(&creation.kind)
.ok_or_else(|| {
ILError::invalid_input(format!("Index kind {} not registered", creation.kind))
})?;
let index_def = Arc::new(IndexDefinition {
index_id,
name: creation.name.clone(),
kind: creation.kind.clone(),
table_id: table.table_id,
table_name: table.table_name.clone(),
table_schema: table.table_schema.clone(),
key_columns: creation.key_columns.clone(),
params: creation.params.clone(),
config: creation.config.clone(),
});
index_kind.supports(&index_def)?;
if let Some(index_id) = tx_helper
.get_index_id(&table.table_id, &creation.name)
.await?
{
return if creation.if_not_exists {
Ok(index_id)
} else {
Err(ILError::invalid_input(format!(
"Index name {} already exists in table {}",
creation.name, table.table_name
)))
};
}
let inline_index_records = build_inline_indexes_for_one_index(
tx_helper,
&table.table_id,
&table.table_schema,
index_kind,
&index_def,
)
.await?;
tx_helper
.insert_inline_indexes(&inline_index_records)
.await?;
let mut projection = vec![0];
for col in creation.key_columns.iter() {
let idx = table.table_schema.arrow_schema.index_of(col)?;
projection.push(idx);
}
projection.sort();
projection.dedup();
let data_file_records = tx_helper.get_data_files(&table.table_id).await?;
let max_concurrency = creation.concurrency.min(data_file_records.len().max(1));
let storage = table.storage.clone();
let table_schema = table.table_schema.clone();
let index_kind = index_kind.clone();
let index_def = index_def.clone();
let namespace_id = table.namespace_id;
let table_id = table.table_id;
let projection = projection.clone();
let results = futures::stream::iter(data_file_records)
.map(|data_file_record| {
let storage = storage.clone();
let table_schema = table_schema.clone();
let index_kind = index_kind.clone();
let index_def = index_def.clone();
let projection = projection.clone();
async move {
let mut index_builder = index_kind.builder(&index_def)?;
let mut stream = read_data_file_by_record(
storage.as_ref(),
&table_schema,
&data_file_record,
Some(projection),
vec![],
1024,
)
.await?;
while let Some(batch) = stream.next().await {
let batch = batch?;
index_builder.append(&batch)?;
}
let index_file_id = Uuid::now_v7();
let relative_path =
IndexFileRecord::build_relative_path(&namespace_id, &table_id, &index_file_id);
let output_file = storage.create(&relative_path).await?;
index_builder.write_file(output_file).await?;
let size = storage.open(&relative_path).await?.metadata().await?.size;
Ok::<_, ILError>(IndexFileRecord {
index_file_id,
table_id,
index_id,
data_file_id: data_file_record.data_file_id,
relative_path,
size: size as i64,
})
}
})
.buffer_unordered(max_concurrency)
.collect::<Vec<_>>()
.await;
let mut index_file_records = Vec::with_capacity(results.len());
for r in results {
index_file_records.push(r?);
}
tx_helper.insert_index_files(&index_file_records).await?;
let key_field_ids = creation
.key_columns
.iter()
.map(|c| Uuid::parse_str(c))
.collect::<Result<Vec<Uuid>, _>>()?;
tx_helper
.insert_index(&IndexRecord {
index_id,
index_name: creation.name.clone(),
index_kind: creation.kind.clone(),
table_id: table.table_id,
key_field_ids,
params: creation.params.encode()?,
config: creation.config.clone(),
})
.await?;
Ok(index_id)
}
pub(crate) async fn build_inline_indexes_for_one_index(
tx_helper: &mut TransactionHelper,
table_id: &Uuid,
table_schema: &TableSchemaRef,
index_kind: &Arc<dyn IndexKind>,
index_def: &IndexDefinitionRef,
) -> ILResult<Vec<InlineIndexRecord>> {
let segment_row_count = index_def.inline_index_segment_row_count(index_kind.as_ref());
let catalog_schema = Arc::new(CatalogSchema::from_arrow(&table_schema.arrow_schema)?);
let row_stream = tx_helper
.scan_inline_rows(table_id, &catalog_schema, &[], None, None)
.await?;
let mut chunk_stream = row_stream.chunks(100);
let mut inline_index_records = Vec::new();
let mut index_builder = index_kind.builder(index_def)?;
let mut counter = 0;
let mut all_row_ids: Vec<Uuid> = Vec::new();
while let Some(row_chunk) = chunk_stream.next().await {
let rows = row_chunk.into_iter().collect::<ILResult<Vec<_>>>()?;
for row in &rows {
if let Some(row_id) = row.get_row_id()? {
all_row_ids.push(row_id);
}
}
counter += rows.len();
let record_batch = rows_to_record_batch(&table_schema.arrow_schema, &rows)?;
index_builder.append(&record_batch)?;
if counter >= segment_row_count && !index_builder.is_empty() {
let mut index_data = Vec::new();
index_builder.write_bytes(&mut index_data)?;
inline_index_records.push(InlineIndexRecord {
index_id: index_builder.index_def().index_id,
created_at: timestamp_ms_from_now(Duration::ZERO),
row_ids: all_row_ids.clone(),
validity: RowValidity::new(all_row_ids.len()),
index_data,
});
tokio::time::sleep(Duration::from_millis(1)).await;
counter = 0;
all_row_ids.clear();
index_builder = index_kind.builder(index_def)?;
}
}
drop(chunk_stream);
if counter > 0 && !index_builder.is_empty() {
let mut index_data = Vec::new();
index_builder.write_bytes(&mut index_data)?;
inline_index_records.push(InlineIndexRecord {
index_id: index_builder.index_def().index_id,
created_at: timestamp_ms_from_now(Duration::ZERO),
row_ids: all_row_ids.clone(),
validity: RowValidity::new(all_row_ids.len()),
index_data,
});
}
Ok(inline_index_records)
}