use super::{
analyzer_registry, Arc, BTreeMap, DocId, Document, Engine, FieldName, FtsIndexStat, SQLError,
TableState, Value,
};
use uqa_storage::InvertedIndex;
type TextIndexDocuments = Vec<(DocId, BTreeMap<FieldName, String>)>;
impl Engine {
pub(crate) fn fts_fields_for_table(&self, name: &str) -> Result<Vec<FieldName>, SQLError> {
Ok(self
.try_query_table(name)
.map_err(|err| SQLError::Internal(format!("resolve table `{name}`: {err}")))?
.map_or_else(Vec::new, |table| table.fts_fields()))
}
pub(crate) fn validate_text_search_field(
&self,
table: &str,
field: &str,
) -> Result<(), SQLError> {
let Some(table_state) = self
.try_query_table(table)
.map_err(|error| SQLError::Internal(format!("resolve text-search table: {error}")))?
else {
return Err(SQLError::UnknownTable(table.to_string()));
};
uqa_sql::semantics::text_indexes::require_physical_text_index(
table,
field,
&table_state.fts_fields(),
|| table_state.columns.read().clone(),
)
}
pub fn fts_index_stats(
&self,
table_filter: Option<&str>,
) -> Result<Vec<FtsIndexStat>, SQLError> {
self.with_direct_read_snapshot(|engine| {
engine.fts_index_stats_with_tables(table_filter, |name| {
engine
.bind_query_table_read(name)
.map(|binding| binding.value)
})
})
}
pub(crate) fn fts_index_stats_in_execution(
&self,
table_filter: Option<&str>,
) -> Result<Vec<FtsIndexStat>, SQLError> {
self.fts_index_stats_with_tables(table_filter, |name| self.require_query_table(name))
}
fn fts_index_stats_with_tables(
&self,
table_filter: Option<&str>,
bind: impl Fn(&str) -> Result<Arc<TableState>, SQLError>,
) -> Result<Vec<FtsIndexStat>, SQLError> {
let mut out = Vec::new();
for table_name in self.fts_stats_table_names(table_filter)? {
let table = bind(&table_name)?;
let mut fields = table.fts_fields();
fields.sort();
let index = table.inverted_index.read();
let index = uqa_execution::serializable::text::ObservedTextIndex::new(
index.as_ref(),
self.serializable_table_state_read(&table)?,
table.columns.snapshot(),
);
for field in fields {
let analyzer = self
.table_field_analyzer_in_execution(&table_name, &field)
.map_err(SQLError::Internal)?
.map_or_else(
|| analyzer_registry::DEFAULT_ANALYZER_NAME.to_string(),
|(name, _)| name,
);
let doc_length_count = index.doc_length_count(Some(&field)).map_err(|error| {
uqa_execution::storage_errors::storage_error(
"read FTS document-length count",
&error,
)
})?;
out.push(FtsIndexStat {
table_name: table_name.clone(),
field: field.clone(),
analyzer,
posting_count: index.posting_count(Some(&field)).map_err(|error| {
uqa_execution::storage_errors::storage_error(
"read FTS posting count",
&error,
)
})?,
doc_length_count,
indexed_doc_count: doc_length_count,
term_count: index.term_count(Some(&field)).map_err(|error| {
uqa_execution::storage_errors::storage_error("read FTS term count", &error)
})?,
total_field_length: index.total_field_length(&field).map_err(|error| {
uqa_execution::storage_errors::storage_error(
"read FTS field length",
&error,
)
})?,
});
}
}
Ok(out)
}
fn fts_stats_table_names(&self, table_filter: Option<&str>) -> Result<Vec<String>, SQLError> {
self.synchronize_table_catalog()
.map_err(|err| SQLError::Internal(format!("refresh table catalog: {err}")))?;
let mut names = if let Some(name) = table_filter {
vec![self
.try_resolve_query_table_name(name)
.map_err(|err| SQLError::Internal(format!("resolve table filter: {err}")))?
.ok_or_else(|| SQLError::UnknownTable(name.to_string()))?]
} else if let Some(tables) = self.query_table_snapshots.as_ref() {
tables
.keys()
.map(uqa_core::RelationIdentity::qualified_name)
.collect()
} else {
self.storage
.tables
.read()
.keys()
.map(uqa_core::RelationIdentity::qualified_name)
.collect()
};
names.sort_unstable();
Ok(names)
}
pub(crate) fn project_fts_sources(t: &Arc<TableState>) -> Result<TextIndexDocuments, String> {
Self::project_fts_sources_inner(t, None)
}
pub(crate) fn project_fts_sources_cancellable(
t: &Arc<TableState>,
cancellation: &uqa_core::CancellationToken,
) -> Result<TextIndexDocuments, String> {
Self::project_fts_sources_inner(t, Some(cancellation))
}
fn project_fts_sources_inner(
t: &Arc<TableState>,
cancellation: Option<&uqa_core::CancellationToken>,
) -> Result<TextIndexDocuments, String> {
if let Some(cancellation) = cancellation {
cancellation.check().map_err(|error| error.to_string())?;
}
let fts_fields = t.fts_fields();
let indexed_docs = {
let store = t.document_store.read();
let doc_ids = store.doc_ids().map_err(|error| error.to_string())?;
let fields: Vec<&str> = fts_fields.iter().map(String::as_str).collect();
let mut indexed_docs = Vec::with_capacity(doc_ids.len());
store
.for_each_fields_multi_ref(&doc_ids, &fields, &mut |doc_id, projected_values| {
if cancellation.is_some_and(uqa_core::CancellationToken::is_cancelled) {
return false;
}
let mut text_fields: BTreeMap<FieldName, String> = BTreeMap::new();
for (field, value) in fts_fields.iter().zip(projected_values) {
if let Value::Str(text) = value {
text_fields.insert(field.clone(), text.clone());
}
}
if !text_fields.is_empty() {
indexed_docs.push((doc_id, text_fields));
}
true
})
.map_err(|error| error.to_string())?;
if let Some(cancellation) = cancellation {
cancellation.check().map_err(|error| error.to_string())?;
}
indexed_docs
};
Ok(indexed_docs)
}
pub(crate) fn rebuild_fts_index(t: &Arc<TableState>) -> Result<(), String> {
let documents = Self::project_fts_sources(t)?;
t.inverted_index
.write()
.try_rebuild_documents(documents)
.map_err(|error| error.to_string())
}
pub(crate) fn rebuild_fts_index_cancellable(
t: &Arc<TableState>,
cancellation: &uqa_core::CancellationToken,
) -> Result<(), String> {
let documents = Self::project_fts_sources_cancellable(t, cancellation)?;
t.inverted_index
.write()
.try_rebuild_documents_cancellable(documents, cancellation)
.map_err(|error| error.to_string())
}
pub fn add_document(
&self,
table: &str,
doc_id: DocId,
document: Document,
) -> Result<(), SQLError> {
self.with_implicit_row_write_transaction(
table,
doc_id,
uqa_sql::ast::LockStrength::ForUpdate,
|engine| engine.add_document_impl(table, doc_id, document, false),
)
}
pub(crate) fn add_document_impl(
&self,
table: &str,
doc_id: DocId,
mut document: Document,
known_new: bool,
) -> Result<(), SQLError> {
uqa_execution::mutation::assignment::refresh_stored_generated_columns(
self.mutation_assignment_context(),
table,
&mut document,
)?;
uqa_execution::serializable::observe_row_write(self, table, doc_id)?;
self.add_prepared_document_impl(table, doc_id, document, known_new)
}
pub(crate) fn add_prepared_document_impl(
&self,
table: &str,
doc_id: DocId,
document: Document,
known_new: bool,
) -> Result<(), SQLError> {
self.add_prepared_document_impl_with_fts(table, doc_id, document, known_new, true, None)
}
pub(crate) fn add_prepared_document_without_fts_impl(
&self,
table: &str,
doc_id: DocId,
document: Document,
known_new: bool,
) -> Result<(), SQLError> {
self.add_prepared_document_impl_with_fts(table, doc_id, document, known_new, false, None)
}
pub(crate) fn add_prepared_stored_document_impl(
&self,
table: &str,
doc_id: DocId,
document: uqa_storage::StoredDocument,
known_new: bool,
) -> Result<(), SQLError> {
let (fields, metadata) = document.into_parts();
self.add_prepared_document_impl_with_fts(
table,
doc_id,
fields,
known_new,
true,
Some(metadata),
)
}
pub(crate) fn prepared_document_text_fields(
&self,
table: &str,
document: &Document,
) -> Result<BTreeMap<FieldName, String>, SQLError> {
let Some(t) = self
.try_table(table)
.map_err(|err| SQLError::Internal(format!("resolve table `{table}`: {err}")))?
else {
return Err(SQLError::UnknownTable(table.to_string()));
};
let mut text_fields = BTreeMap::new();
for name in &t.fts_fields() {
if let Some(Value::Str(value)) = document.get(name) {
text_fields.insert(name.clone(), value.clone());
}
}
Ok(text_fields)
}
pub(crate) fn add_prepared_fts_documents(
&self,
table: &str,
documents: Vec<(DocId, BTreeMap<FieldName, String>)>,
) -> Result<(), SQLError> {
let Some(t) = self
.try_table(table)
.map_err(|err| SQLError::Internal(format!("resolve table `{table}`: {err}")))?
else {
return Err(SQLError::UnknownTable(table.to_string()));
};
let result = uqa_execution::serializable::text::add_documents(
self,
table,
t.columns.snapshot(),
t.inverted_index.write().as_mut(),
documents,
);
result
}
fn add_prepared_document_impl_with_fts(
&self,
table: &str,
doc_id: DocId,
mut document: Document,
known_new: bool,
index_fts: bool,
metadata: Option<uqa_storage::DocumentMetadata>,
) -> Result<(), SQLError> {
let Some(table_name) = self
.try_resolve_table_name(table)
.map_err(|err| SQLError::Internal(format!("resolve table `{table}`: {err}")))?
else {
return Err(SQLError::UnknownTable(table.to_string()));
};
let Some(t) = self
.try_table(table)
.map_err(|err| SQLError::Internal(format!("resolve table `{table}`: {err}")))?
else {
return Err(SQLError::UnknownTable(table.to_string()));
};
let existed = if known_new {
false
} else {
t.document_store
.read()
.get(doc_id)
.map_err(|error| SQLError::Internal(format!("read existing document: {error}")))?
.is_some()
};
let (old_indexed, indexed_fields) = if known_new {
(None, Self::value_indexes_built_fields(&t))
} else {
let old = Self::value_indexes_old_values(&t, doc_id);
let fields = old.as_ref().map(|old| {
old.keys()
.cloned()
.collect::<Vec<uqa_storage::ValueIndexKey>>()
});
(old, fields)
};
let persistent_indexed =
self.persistent_value_index_document_values(&table_name, &document)?;
let new_indexed = indexed_fields
.map(|fields| {
if let Some(values) = &persistent_indexed {
return Ok(values.clone());
}
self.value_index_document_values(&table_name, &fields, &document)
})
.transpose()?;
self.observe_value_index_write(
&table_name,
&t,
doc_id,
existed,
old_indexed.as_ref(),
persistent_indexed.as_ref().or(new_indexed.as_ref()),
)?;
if index_fts {
let text_fields = self.prepared_document_text_fields(table, &document)?;
uqa_execution::serializable::text::add_document(
self,
&table_name,
t.columns.snapshot(),
t.inverted_index.write().as_mut(),
doc_id,
text_fields,
)?;
}
let columns = t.columns.read().clone();
crate::generated::strip_virtual_generated_columns(&columns, &mut document);
let metadata = match metadata {
Some(metadata) => metadata,
None => uqa_storage::DocumentMetadata::with_tuple_xmin(self.tuple_version_xid()?),
};
self.advance_next_id(&table_name, doc_id).map_err(|error| {
uqa_execution::mutation::errors::identifier_storage_error(
"observe inserted document identity",
&error,
)
})?;
let mut store = t.document_store.write();
store
.put_stored(
doc_id,
uqa_storage::StoredDocument::with_metadata(document, metadata),
)
.map_err(|err| crate::table_storage::document_store_write_error(&err))?;
if let Some(new) = persistent_indexed.as_ref() {
self.persist_value_indexes_apply_write(&table_name, doc_id, Some(new))?;
}
if let Some(new) = new_indexed.as_ref() {
Self::value_indexes_apply_write(&t, doc_id, old_indexed.as_ref(), Some(new));
}
drop(store);
self.mark_column_stats_dirty(&table_name, &t)
.map_err(|err| SQLError::Internal(format!("invalidate column stats: {err}")))?;
if existed {
self.note_row_changed(&table_name, doc_id)?;
} else {
self.note_row_inserted(&table_name, doc_id)?;
}
Ok(())
}
}