From 2b6856035910018bf349f1ceddb8756596364324 Mon Sep 17 00:00:00 2001 From: Priec Date: Thu, 13 Aug 2026 22:21:31 +0200 Subject: [PATCH] search fix to new architecture3 --- common/src/search.rs | 4 +- search/src/lib.rs | 464 ++++++++++++++++++++++++++++++++----------- server | 2 +- 3 files changed, 347 insertions(+), 123 deletions(-) diff --git a/common/src/search.rs b/common/src/search.rs index 935d361e..8c7f0309 100644 --- a/common/src/search.rs +++ b/common/src/search.rs @@ -23,7 +23,7 @@ pub const F_DATA_NGRAM: &str = "data_ngram"; pub const F_DATA_EXACT: &str = "data_exact"; pub const JOURNAL_TABLE_NAME: &str = "general_ledger"; pub const ARCHIVED_ROW_KEY_PREFIX: &str = "__archive__:"; -pub const SEARCH_INDEX_FORMAT_DIRECTORY: &str = "v3"; +pub const SEARCH_INDEX_FORMAT_DIRECTORY: &str = "v4"; /// Root for the current on-disk index format. The format component prevents /// an executable with new projection semantics from opening old documents. @@ -134,7 +134,7 @@ pub fn create_search_schema() -> Schema { schema_builder.add_u64_field(F_IS_ARCHIVED, INDEXED | STORED | FAST); schema_builder.add_u64_field(F_TABLE_DEFINITION_ID, STORED | FAST); schema_builder.add_u64_field(F_VERSION_NUMBER, STORED | FAST); - schema_builder.add_text_field(F_TABLE_NAME, STRING | STORED); + schema_builder.add_text_field(F_TABLE_NAME, STRING | STORED | FAST); schema_builder.add_text_field(F_ROW_KEY, STRING | STORED); schema_builder.add_text_field(F_ALL_TEXT, text_options(TOK_WORD)); diff --git a/search/src/lib.rs b/search/src/lib.rs index c9dbd919..e28f8cf7 100644 --- a/search/src/lib.rs +++ b/search/src/lib.rs @@ -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> = 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 { - 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, + doc_id, + field_name: &str| + -> Result { + 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) { + 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, + revisions: tantivy::columnar::Column, + archived_flags: tantivy::columnar::Column, + definition_ids: tantivy::columnar::Column, + versions: tantivy::columnar::Column, + table_names: tantivy::columnar::StrColumn, +} + +fn candidate_segment_fields( + searcher: &tantivy::Searcher, + segment_ord: u32, +) -> Result { + 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, 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, 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)>, 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::>(); + let revisions = rows.iter().map(|item| item.1).collect::>(); + 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::>(); + let row_ids = archives.iter().map(|item| item.1).collect::>(); + let versions = archives.iter().map(|item| item.2).collect::>(); + 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, 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::>(); - 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 = 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::>(); + 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::>(); - if order.is_some() { - Ok(hits) - } else { - let page = hits.into_iter().skip(offset).take(limit).collect::>(); - 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( diff --git a/server b/server index 714104f7..d0504ad2 160000 --- a/server +++ b/server @@ -1 +1 @@ -Subproject commit 714104f7605806bf407337fa639418f350d15e60 +Subproject commit d0504ad2a963671674eae1ec93e03bc5fc0c0f1b