filipriec/komp_ac
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

search fix to new architecture3

Browse Source
This commit is contained in:
Priec
2026-08-13 22:21:31 +02:00
parent 09b4d5097a
commit 2b68560359
3 changed files with 347 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

464
search/src/lib.rs
View File

@@ -11,8 +11,8 @@ use common::proto::komp_ac::search::{
search_response::Hit,
};
use common::search::{
F_IS_ARCHIVED, F_PG_ID, F_ROW_REVISION, F_TABLE_DEFINITION_ID, F_VERSION_NUMBER,
SchemaFields, canonical_exact_search_value, parse_archived_search_row_key,
F_IS_ARCHIVED, F_PG_ID, F_ROW_REVISION, F_TABLE_DEFINITION_ID, F_TABLE_NAME,
F_VERSION_NUMBER, SchemaFields, canonical_exact_search_value, parse_archived_search_row_key,
register_tokenizers, search_index_path, search_row_key,
};
use common::system_column::{internal_column_names, is_internal_column, is_system_column};
@@ -22,8 +22,7 @@ use query_builder::{
use sqlx::{AssertSqlSafe, PgPool, Row};
use tantivy::collector::TopDocs;
use tantivy::query::Query;
use tantivy::schema::Value;
use tantivy::{Index, IndexReader, ReloadPolicy, TantivyDocument};
use tantivy::{Index, IndexReader, ReloadPolicy};
use tonic::{Request, Response, Status};
use tracing::info;
@@ -32,7 +31,6 @@ const HARD_RESULT_LIMIT: usize = 200;
const DEFAULT_LIST_LIMIT: usize = 5;
const SEARCH_SCORE_RELATIVE_FLOOR: f32 = 0.25;
const SEARCH_SCORE_GROUP_WIDTH: f32 = 0.25;
const SEARCH_VALIDATION_OVERFETCH_FACTOR: usize = 4;
pub struct SearcherService {
pub pool: PgPool,
@@ -125,7 +123,6 @@ impl SearcherService {
normalized.offset,
normalized.order.as_ref(),
normalized.version_scope,
false,
)
.await?;
@@ -224,35 +221,67 @@ async fn count_authoritative_matches(
} else {
0.0
};
let mut documents_by_segment: HashMap<u32, Vec<u32>> = HashMap::new();
for (score, address) in documents {
if score >= score_floor {
documents_by_segment
.entry(address.segment_ord)
.or_default()
.push(address.doc_id);
}
}
let mut current_rows = Vec::new();
let mut archives = Vec::new();
for (score, address) in documents {
if score < score_floor {
continue;
}
let segment = searcher.segment_reader(address.segment_ord);
let fast_u64 = |field_name: &str| -> Result<i64, Status> {
segment
.fast_fields()
.u64(field_name)
.map_err(|error| Status::internal(format!(
"Search count fast field '{field_name}' failed: {error}"
)))?
.first(address.doc_id)
for (segment_ord, segment_documents) in documents_by_segment {
let segment = searcher.segment_reader(segment_ord);
let fast_fields = segment.fast_fields();
let pg_ids = fast_fields.u64(F_PG_ID).map_err(|error| {
Status::internal(format!("Search count fast field '{F_PG_ID}' failed: {error}"))
})?;
let archived_flags = fast_fields.u64(F_IS_ARCHIVED).map_err(|error| {
Status::internal(format!(
"Search count fast field '{F_IS_ARCHIVED}' failed: {error}"
))
})?;
let revisions = fast_fields.u64(F_ROW_REVISION).map_err(|error| {
Status::internal(format!(
"Search count fast field '{F_ROW_REVISION}' failed: {error}"
))
})?;
let definition_ids = fast_fields.u64(F_TABLE_DEFINITION_ID).map_err(|error| {
Status::internal(format!(
"Search count fast field '{F_TABLE_DEFINITION_ID}' failed: {error}"
))
})?;
let versions = fast_fields.u64(F_VERSION_NUMBER).map_err(|error| {
Status::internal(format!(
"Search count fast field '{F_VERSION_NUMBER}' failed: {error}"
))
})?;
let read = |column: &tantivy::columnar::Column<u64>,
doc_id,
field_name: &str|
-> Result<i64, Status> {
column
.first(doc_id)
.and_then(|value| i64::try_from(value).ok())
.ok_or_else(|| Status::internal(format!(
"Search count document has no valid '{field_name}'"
)))
.ok_or_else(|| {
Status::internal(format!(
"Search count document has no valid '{field_name}'"
))
})
};
let row_id = fast_u64(F_PG_ID)?;
if fast_u64(F_IS_ARCHIVED)? != 0 {
archives.push((
fast_u64(F_TABLE_DEFINITION_ID)?,
row_id,
fast_u64(F_VERSION_NUMBER)?,
));
} else {
current_rows.push((row_id, fast_u64(F_ROW_REVISION)?));
for doc_id in segment_documents {
let row_id = read(&pg_ids, doc_id, F_PG_ID)?;
if read(&archived_flags, doc_id, F_IS_ARCHIVED)? != 0 {
archives.push((
read(&definition_ids, doc_id, F_TABLE_DEFINITION_ID)?,
row_id,
read(&versions, doc_id, F_VERSION_NUMBER)?,
));
} else {
current_rows.push((row_id, read(&revisions, doc_id, F_ROW_REVISION)?));
}
}
}
current_rows.sort_unstable();
@@ -425,6 +454,131 @@ struct SearchCandidate {
row_key: String,
}
impl SearchCandidate {
fn validation_key(&self) -> (String, Option<i64>) {
if parse_archived_search_row_key(&self.row_key).is_some() {
(self.row_key.clone(), None)
} else {
(self.row_key.clone(), Some(self.row_revision))
}
}
}
struct CandidateSegmentFields {
pg_ids: tantivy::columnar::Column<u64>,
revisions: tantivy::columnar::Column<u64>,
archived_flags: tantivy::columnar::Column<u64>,
definition_ids: tantivy::columnar::Column<u64>,
versions: tantivy::columnar::Column<u64>,
table_names: tantivy::columnar::StrColumn,
}
fn candidate_segment_fields(
searcher: &tantivy::Searcher,
segment_ord: u32,
) -> Result<CandidateSegmentFields, Status> {
let segment = searcher.segment_reader(segment_ord);
let fast_fields = segment.fast_fields();
let u64_column = |field_name: &str| {
fast_fields.u64(field_name).map_err(|error| {
Status::internal(format!(
"Search candidate fast field '{field_name}' failed: {error}"
))
})
};
let table_names = fast_fields
.str(F_TABLE_NAME)
.map_err(|error| {
Status::internal(format!(
"Search candidate fast field '{F_TABLE_NAME}' failed: {error}"
))
})?
.ok_or_else(|| {
Status::internal(format!(
"Search candidate fast field '{F_TABLE_NAME}' is missing"
))
})?;
Ok(CandidateSegmentFields {
pg_ids: u64_column(F_PG_ID)?,
revisions: u64_column(F_ROW_REVISION)?,
archived_flags: u64_column(F_IS_ARCHIVED)?,
definition_ids: u64_column(F_TABLE_DEFINITION_ID)?,
versions: u64_column(F_VERSION_NUMBER)?,
table_names,
})
}
fn search_candidates_from_documents(
searcher: &tantivy::Searcher,
documents: Vec<(f32, tantivy::DocAddress)>,
) -> Result<Vec<SearchCandidate>, Status> {
let mut fields_by_segment = HashMap::new();
for (_, address) in &documents {
if !fields_by_segment.contains_key(&address.segment_ord) {
fields_by_segment.insert(
address.segment_ord,
candidate_segment_fields(searcher, address.segment_ord)?,
);
}
}
let mut candidates = Vec::with_capacity(documents.len());
for (score, address) in documents {
let fields = fields_by_segment.get(&address.segment_ord).ok_or_else(|| {
Status::internal("Search candidate segment fields were not loaded")
})?;
let read_u64 = |column: &tantivy::columnar::Column<u64>, field_name: &str| {
column.first(address.doc_id).ok_or_else(|| {
Status::internal(format!(
"Search candidate document has no '{field_name}'"
))
})
};
let row_id = i64::try_from(read_u64(&fields.pg_ids, F_PG_ID)?)
.map_err(|_| Status::internal("Search candidate row id is out of range"))?;
let row_revision = i64::try_from(read_u64(&fields.revisions, F_ROW_REVISION)?)
.map_err(|_| Status::internal("Search candidate revision is out of range"))?;
let archived = read_u64(&fields.archived_flags, F_IS_ARCHIVED)? != 0;
let table_name_ord = fields
.table_names
.ords()
.first(address.doc_id)
.ok_or_else(|| Status::internal("Search candidate document has no table name"))?;
let mut table_name = String::new();
if !fields
.table_names
.ord_to_str(table_name_ord, &mut table_name)
.map_err(|error| {
Status::internal(format!("Search candidate table name read failed: {error}"))
})?
{
return Err(Status::internal(
"Search candidate table name dictionary entry is missing",
));
}
let row_key = if archived {
let definition_id = i64::try_from(read_u64(
&fields.definition_ids,
F_TABLE_DEFINITION_ID,
)?)
.map_err(|_| Status::internal("Search candidate table id is out of range"))?;
let version = i64::try_from(read_u64(&fields.versions, F_VERSION_NUMBER)?)
.map_err(|_| Status::internal("Search candidate version is out of range"))?;
common::search::archived_search_row_key(definition_id, row_id, version)
} else {
search_row_key(&table_name, row_id)
};
candidates.push(SearchCandidate {
score,
row_id,
row_revision,
table_name,
row_key,
});
}
Ok(candidates)
}
impl NormalizedSearchRequest {
fn has_input(&self) -> bool {
!self.free_query.is_empty() || !self.must.is_empty()
@@ -1153,6 +1307,110 @@ async fn fetch_ordered_rows(
.collect()
}
async fn validate_search_candidates(
pool: &PgPool,
profile_name: &str,
candidates: &[SearchCandidate],
) -> Result<HashSet<(String, Option<i64>)>, Status> {
let mut valid_keys = HashSet::new();
let mut current_by_table: HashMap<&str, Vec<(i64, i64)>> = HashMap::new();
let mut archives_by_table: HashMap<&str, Vec<(i64, i64, i64)>> = HashMap::new();
for candidate in candidates {
if let Some(archive) = parse_archived_search_row_key(&candidate.row_key) {
archives_by_table
.entry(&candidate.table_name)
.or_default()
.push(archive);
} else {
current_by_table
.entry(&candidate.table_name)
.or_default()
.push((candidate.row_id, candidate.row_revision));
}
}
for (table_name, rows) in current_by_table {
let qualified_table = qualified_visible_table(pool, profile_name, table_name).await?;
let ids = rows.iter().map(|item| item.0).collect::<Vec<_>>();
let revisions = rows.iter().map(|item| item.1).collect::<Vec<_>>();
let sql = format!(
"SELECT candidate.id, candidate.row_revision \
FROM UNNEST($1::BIGINT[], $2::BIGINT[]) \
AS candidate(id, row_revision) \
JOIN {qualified_table} current_row \
ON current_row.id = candidate.id \
AND current_row.row_revision = candidate.row_revision \
WHERE current_row.deleted = FALSE"
);
let valid_rows = sqlx::query(AssertSqlSafe(sql))
.bind(&ids)
.bind(&revisions)
.fetch_all(pool)
.await
.map_err(|error| {
Status::internal(format!("Current search candidate validation failed: {error}"))
})?;
for row in valid_rows {
let row_id: i64 = row.try_get("id").map_err(|error| {
Status::internal(format!("Current search id read failed: {error}"))
})?;
let revision: i64 = row.try_get("row_revision").map_err(|error| {
Status::internal(format!("Current search revision read failed: {error}"))
})?;
valid_keys.insert((search_row_key(table_name, row_id), Some(revision)));
}
}
for (table_name, archives) in archives_by_table {
let definition_ids = archives.iter().map(|item| item.0).collect::<Vec<_>>();
let row_ids = archives.iter().map(|item| item.1).collect::<Vec<_>>();
let versions = archives.iter().map(|item| item.2).collect::<Vec<_>>();
let rows = sqlx::query(
r#"SELECT archive.table_definition_id, archive.source_row_id,
archive.version_number
FROM UNNEST($1::BIGINT[], $2::BIGINT[], $3::BIGINT[])
AS requested(table_definition_id, source_row_id, version_number)
JOIN table_row_archives archive
ON archive.table_definition_id = requested.table_definition_id
AND archive.source_row_id = requested.source_row_id
AND archive.version_number = requested.version_number
JOIN table_definitions definition ON definition.id = archive.table_definition_id
JOIN schemas owner ON owner.id = definition.schema_id
WHERE definition.table_name = $4
AND definition.deleted = FALSE
AND (owner.name = $5 OR definition.is_global = TRUE)"#,
)
.bind(&definition_ids)
.bind(&row_ids)
.bind(&versions)
.bind(table_name)
.bind(profile_name)
.fetch_all(pool)
.await
.map_err(|error| {
Status::internal(format!("Archived search candidate validation failed: {error}"))
})?;
for row in rows {
let definition_id: i64 = row.try_get("table_definition_id").map_err(|error| {
Status::internal(format!("Archived table id read failed: {error}"))
})?;
let row_id: i64 = row.try_get("source_row_id").map_err(|error| {
Status::internal(format!("Archived row id read failed: {error}"))
})?;
let version: i64 = row.try_get("version_number").map_err(|error| {
Status::internal(format!("Archived version read failed: {error}"))
})?;
valid_keys.insert((
common::search::archived_search_row_key(definition_id, row_id, version),
None,
));
}
}
Ok(valid_keys)
}
async fn run_search(
pool: &PgPool,
profile: &ProfileIndex,
@@ -1164,7 +1422,6 @@ async fn run_search(
offset: usize,
order: Option<&NormalizedSearchOrder>,
version_scope: SearchVersionScope,
exhaustive_validation: bool,
) -> Result<Vec<Hit>, Status> {
let master_query = build_master_query(
&profile.index,
@@ -1176,78 +1433,69 @@ async fn run_search(
)?;
let searcher = profile.reader.searcher();
let window_limit = if order.is_some() || exhaustive_validation {
searcher.num_docs() as usize
} else {
offset
.saturating_add(limit)
.saturating_mul(SEARCH_VALIDATION_OVERFETCH_FACTOR)
.min(searcher.num_docs() as usize)
};
if window_limit == 0 {
let num_docs = searcher.num_docs() as usize;
if num_docs == 0 {
return Ok(Vec::new());
}
let top_docs = searcher
.search(&*master_query, &TopDocs::with_limit(window_limit).order_by_score())
.map_err(|e| Status::internal(format!("Search failed: {}", e)))?;
if top_docs.is_empty() {
return Ok(vec![]);
}
let best_score = top_docs.first().map(|(score, _)| *score).unwrap_or_default();
let score_floor = if best_score > 0.0 {
best_score * SEARCH_SCORE_RELATIVE_FLOOR
let requested_candidates = offset.saturating_add(limit).max(1);
let mut window_limit = if order.is_some() {
num_docs
} else {
0.0
requested_candidates.min(num_docs)
};
let eligible_docs = top_docs
.into_iter()
.filter(|(score, _)| *score >= score_floor)
.collect::<Vec<_>>();
let page_docs = eligible_docs;
let candidates = loop {
let top_docs = searcher
.search(
&*master_query,
&TopDocs::with_limit(window_limit).order_by_score(),
)
.map_err(|e| Status::internal(format!("Search failed: {}", e)))?;
if top_docs.is_empty() {
return Ok(Vec::new());
}
let mut candidates: Vec<SearchCandidate> = Vec::with_capacity(page_docs.len());
for (score, doc_address) in page_docs {
let doc: TantivyDocument = searcher
.doc(doc_address)
.map_err(|e| Status::internal(format!("Failed to retrieve document: {}", e)))?;
let Some(pg_id) = doc
.get_first(profile.fields.pg_id)
.and_then(|value| value.as_u64())
else {
continue;
let collected_docs = top_docs.len();
let best_score = top_docs.first().map(|(score, _)| *score).unwrap_or_default();
let score_floor = if best_score > 0.0 {
best_score * SEARCH_SCORE_RELATIVE_FLOOR
} else {
0.0
};
let Some(table_name) = doc
.get_first(profile.fields.table_name)
.and_then(|value| value.as_str())
else {
continue;
};
let Some(row_key) = doc
.get_first(profile.fields.row_key)
.and_then(|value| value.as_str())
else {
continue;
};
let Some(row_revision) = doc
.get_first(profile.fields.row_revision)
.and_then(|value| value.as_u64())
.and_then(|value| i64::try_from(value).ok())
else {
continue;
};
candidates.push(SearchCandidate {
score,
row_id: pg_id as i64,
row_revision,
table_name: table_name.to_string(),
row_key: row_key.to_string(),
});
}
let reached_score_floor = top_docs
.last()
.is_some_and(|(score, _)| *score < score_floor);
let eligible_docs = top_docs
.into_iter()
.filter(|(score, _)| *score >= score_floor)
.collect::<Vec<_>>();
let candidates = search_candidates_from_documents(&searcher, eligible_docs)?;
if order.is_some() {
break candidates;
}
let valid_keys = validate_search_candidates(pool, profile_name, &candidates).await?;
let valid_count = candidates
.iter()
.filter(|candidate| valid_keys.contains(&candidate.validation_key()))
.count();
if valid_count >= requested_candidates
|| window_limit == num_docs
|| reached_score_floor
|| collected_docs < window_limit
{
break candidates
.into_iter()
.filter(|candidate| valid_keys.contains(&candidate.validation_key()))
.skip(offset)
.take(limit)
.collect();
}
window_limit = window_limit.saturating_mul(2).min(num_docs);
};
if candidates.is_empty() {
return Ok(vec![]);
return Ok(Vec::new());
}
if let Some(order) = order {
@@ -1421,31 +1669,7 @@ async fn run_search(
})
})
.collect::<Vec<_>>();
if order.is_some() {
Ok(hits)
} else {
let page = hits.into_iter().skip(offset).take(limit).collect::<Vec<_>>();
if !exhaustive_validation
&& page.len() < limit
&& window_limit < searcher.num_docs() as usize
{
return Box::pin(run_search(
pool,
profile,
profile_name,
table_filter,
free_query,
must,
limit,
offset,
order,
version_scope,
true,
))
.await;
}
Ok(page)
}
Ok(hits)
}
async fn fetch_ordered_candidate_rows(