search fix to new architecture3

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));
 

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,
+    F_IS_ARCHIVED, F_PG_ID, F_ROW_REVISION, F_TABLE_DEFINITION_ID, F_TABLE_NAME,
-    SchemaFields, canonical_exact_search_value, parse_archived_search_row_key,
+    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};
-use tantivy::{Index, IndexReader, ReloadPolicy, TantivyDocument};
 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 {
+    for (segment_ord, segment_documents) in documents_by_segment {
-        if score < score_floor {
+        let segment = searcher.segment_reader(segment_ord);
-            continue;
+        let fast_fields = segment.fast_fields();
-        }
+        let pg_ids = fast_fields.u64(F_PG_ID).map_err(|error| {
-        let segment = searcher.segment_reader(address.segment_ord);
+            Status::internal(format!("Search count fast field '{F_PG_ID}' failed: {error}"))
-        let fast_u64 = |field_name: &str| -> Result<i64, Status> {
+        })?;
-            segment
+        let archived_flags = fast_fields.u64(F_IS_ARCHIVED).map_err(|error| {
-                .fast_fields()
+            Status::internal(format!(
-                .u64(field_name)
+                "Search count fast field '{F_IS_ARCHIVED}' failed: {error}"
-                .map_err(|error| Status::internal(format!(
+            ))
-                    "Search count fast field '{field_name}' failed: {error}"
+        })?;
-                )))?
+        let revisions = fast_fields.u64(F_ROW_REVISION).map_err(|error| {
-                .first(address.doc_id)
+            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!(
+                .ok_or_else(|| {
+                    Status::internal(format!(
                         "Search count document has no valid '{field_name}'"
-                )))
+                    ))
+                })
         };
-        let row_id = fast_u64(F_PG_ID)?;
+        for doc_id in segment_documents {
-        if fast_u64(F_IS_ARCHIVED)? != 0 {
+            let row_id = read(&pg_ids, doc_id, F_PG_ID)?;
+            if read(&archived_flags, doc_id, F_IS_ARCHIVED)? != 0 {
                 archives.push((
-                fast_u64(F_TABLE_DEFINITION_ID)?,
+                    read(&definition_ids, doc_id, F_TABLE_DEFINITION_ID)?,
                     row_id,
-                fast_u64(F_VERSION_NUMBER)?,
+                    read(&versions, doc_id, F_VERSION_NUMBER)?,
                 ));
             } else {
-            current_rows.push((row_id, fast_u64(F_ROW_REVISION)?));
+                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 {
+    let num_docs = searcher.num_docs() as usize;
-        searcher.num_docs() as usize
+    if num_docs == 0 {
-    } else {
-        offset
-            .saturating_add(limit)
-            .saturating_mul(SEARCH_VALIDATION_OVERFETCH_FACTOR)
-            .min(searcher.num_docs() as usize)
-    };
-    if window_limit == 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)))?;
 
+    let requested_candidates = offset.saturating_add(limit).max(1);
+    let mut window_limit = if order.is_some() {
+        num_docs
+    } else {
+        requested_candidates.min(num_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![]);
+            return Ok(Vec::new());
         }
 
+        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 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 page_docs = eligible_docs;
+        let candidates = search_candidates_from_documents(&searcher, eligible_docs)?;
 
-    let mut candidates: Vec<SearchCandidate> = Vec::with_capacity(page_docs.len());
+        if order.is_some() {
-    for (score, doc_address) in page_docs {
+            break candidates;
-        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 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 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)
-    }
 }
 
 async fn fetch_ordered_candidate_rows(