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

search fix to new architecture2

Browse Source
This commit is contained in:
Priec
2026-08-13 21:57:27 +02:00
parent 810c8ffd08
commit 09b4d5097a
7 changed files with 251 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
search/src/lib.rs
View File

@@ -11,8 +11,9 @@ use common::proto::komp_ac::search::{
search_response::Hit,
};
use common::search::{
SchemaFields, parse_archived_search_row_key, register_tokenizers, search_index_path,
search_row_key,
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,
register_tokenizers, search_index_path, search_row_key,
};
use common::system_column::{internal_column_names, is_internal_column, is_system_column};
use query_builder::{
@@ -31,6 +32,7 @@ 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,
@@ -112,10 +114,6 @@ impl SearcherService {
.await?;
let profile = profile_index(&self.profiles, &normalized.profile_name, &index_path)?;
profile
.reader
.reload()
.map_err(|error| Status::internal(format!("Search index reload failed: {error}")))?;
let mut hits = run_search(
&self.pool,
&profile,
@@ -127,6 +125,7 @@ impl SearcherService {
normalized.offset,
normalized.order.as_ref(),
normalized.version_scope,
false,
)
.await?;
@@ -187,10 +186,6 @@ impl SearcherService {
)
.await?;
let profile = profile_index(&self.profiles, &normalized.profile_name, &index_path)?;
profile
.reader
.reload()
.map_err(|error| Status::internal(format!("Search index reload failed: {error}")))?;
let query = build_master_query(
&profile.index,
&profile.fields,
@@ -219,12 +214,7 @@ async fn count_authoritative_matches(
query: &dyn Query,
) -> Result<u64, Status> {
let searcher = profile.reader.searcher();
let documents = searcher
.search(
query,
&TopDocs::with_limit(searcher.num_docs() as usize).order_by_score(),
)
.map_err(|error| Status::internal(format!("Search count failed: {error}")))?;
let documents = top_documents_for_count(&searcher, query)?;
let best_score = documents
.first()
.map(|(score, _)| *score)
@@ -234,44 +224,92 @@ async fn count_authoritative_matches(
} else {
0.0
};
let mut current_ids = Vec::new();
let mut current_rows = Vec::new();
let mut archives = Vec::new();
for (score, address) in documents {
if score < score_floor {
continue;
}
let document: TantivyDocument = searcher
.doc(address)
.map_err(|error| Status::internal(format!("Search count document read failed: {error}")))?;
let Some(row_key) = document
.get_first(profile.fields.row_key)
.and_then(|value| value.as_str())
else {
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)
.and_then(|value| i64::try_from(value).ok())
.ok_or_else(|| Status::internal(format!(
"Search count document has no valid '{field_name}'"
)))
};
if let Some(archive) = parse_archived_search_row_key(row_key) {
archives.push(archive);
} else if let Some(row_id) = document
.get_first(profile.fields.pg_id)
.and_then(|value| value.as_u64())
.and_then(|value| i64::try_from(value).ok())
{
current_ids.push(row_id);
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)?));
}
}
current_ids.sort_unstable();
current_ids.dedup();
current_rows.sort_unstable();
current_rows.dedup();
archives.sort_unstable();
archives.dedup();
let current_count = if current_ids.is_empty() {
count_validated_candidates(
pool,
profile_name,
table_name,
&current_rows,
&archives,
)
.await
}
fn top_documents_for_count(
searcher: &tantivy::Searcher,
query: &dyn Query,
) -> Result<Vec<(f32, tantivy::DocAddress)>, Status> {
if searcher.num_docs() == 0 {
return Ok(Vec::new());
}
// DocAddress is local to this exact Searcher snapshot. The caller must use
// the same searcher to load every address returned here.
searcher
.search(
query,
&TopDocs::with_limit(searcher.num_docs() as usize).order_by_score(),
)
.map_err(|error| Status::internal(format!("Search count failed: {error}")))
}
async fn count_validated_candidates(
pool: &PgPool,
profile_name: &str,
table_name: &str,
current_rows: &[(i64, i64)],
archives: &[(i64, i64, i64)],
) -> Result<u64, Status> {
let current_count = if current_rows.is_empty() {
0
} else {
let qualified_table = qualified_visible_table(pool, profile_name, table_name).await?;
let current_ids = current_rows.iter().map(|item| item.0).collect::<Vec<_>>();
let current_revisions = current_rows.iter().map(|item| item.1).collect::<Vec<_>>();
sqlx::query_scalar::<_, i64>(AssertSqlSafe(format!(
"SELECT COUNT(*) FROM {qualified_table} WHERE deleted = FALSE AND id = ANY($1)"
"SELECT COUNT(*) FROM {qualified_table} current_row \
JOIN UNNEST($1::BIGINT[], $2::BIGINT[]) \
AS candidate(id, row_revision) \
ON candidate.id = current_row.id \
AND candidate.row_revision = current_row.row_revision \
WHERE current_row.deleted = FALSE"
)))
.bind(&current_ids)
.bind(&current_revisions)
.fetch_one(pool)
.await
.map_err(|error| Status::internal(format!("Current search count validation failed: {error}")))?
@@ -367,26 +405,7 @@ struct NormalizedColumnConstraint {
}
fn normalize_constraint_value(query: &str, field_type: &str) -> Result<String, Status> {
let normalized_type = field_type.trim().to_ascii_lowercase();
if normalized_type == "money" || normalized_type.starts_with("decimal(") {
return query
.parse::<rust_decimal::Decimal>()
.map(|value| value.normalize().to_string())
.map_err(|error| Status::invalid_argument(format!(
"Exact numeric search value is invalid: {error}"
)));
}
if matches!(normalized_type.as_str(), "integer" | "bigint")
|| normalized_type.starts_with("link(")
{
return query
.parse::<i64>()
.map(|value| value.to_string())
.map_err(|error| Status::invalid_argument(format!(
"Exact integer search value is invalid: {error}"
)));
}
Ok(query.to_string())
canonical_exact_search_value(query, field_type).map_err(Status::invalid_argument)
}
fn public_search_system_type(column: &str) -> Option<&'static str> {
@@ -401,6 +420,7 @@ fn public_search_system_type(column: &str) -> Option<&'static str> {
struct SearchCandidate {
score: f32,
row_id: i64,
row_revision: i64,
table_name: String,
row_key: String,
}
@@ -1144,6 +1164,7 @@ 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,
@@ -1155,10 +1176,13 @@ async fn run_search(
)?;
let searcher = profile.reader.searcher();
let window_limit = if order.is_some() {
let window_limit = if order.is_some() || exhaustive_validation {
searcher.num_docs() as usize
} else {
offset.saturating_add(limit)
offset
.saturating_add(limit)
.saturating_mul(SEARCH_VALIDATION_OVERFETCH_FACTOR)
.min(searcher.num_docs() as usize)
};
if window_limit == 0 {
return Ok(Vec::new());
@@ -1181,15 +1205,7 @@ async fn run_search(
.into_iter()
.filter(|(score, _)| *score >= score_floor)
.collect::<Vec<_>>();
let page_docs = if order.is_some() {
eligible_docs
} else {
eligible_docs
.into_iter()
.skip(offset)
.take(limit)
.collect::<Vec<_>>()
};
let page_docs = eligible_docs;
let mut candidates: Vec<SearchCandidate> = Vec::with_capacity(page_docs.len());
for (score, doc_address) in page_docs {
@@ -1214,9 +1230,17 @@ async fn run_search(
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(),
});
@@ -1240,20 +1264,20 @@ async fn run_search(
.await;
}
let mut ids_by_table: HashMap<String, Vec<i64>> = HashMap::new();
let mut rows_by_table: HashMap<String, Vec<(i64, i64)>> = HashMap::new();
for candidate in &candidates {
if parse_archived_search_row_key(&candidate.row_key).is_some() {
continue;
}
ids_by_table
rows_by_table
.entry(candidate.table_name.clone())
.or_default()
.push(candidate.row_id);
.push((candidate.row_id, candidate.row_revision));
}
let mut content_map: HashMap<String, (String, Vec<String>, Vec<String>, i64, bool)> =
HashMap::new();
for (table_name, pg_ids) in ids_by_table {
for (table_name, candidate_rows) in rows_by_table {
validate_identifier(&table_name, "table_name")?;
let physical_to_display =
table_physical_to_display_map(pool, profile_name, &table_name).await?;
@@ -1261,6 +1285,8 @@ async fn run_search(
table_internal_column_names(pool, profile_name, &table_name).await?;
let display_columns = table_row_display_columns(pool, profile_name, &table_name).await?;
let qualified_table = qualified_visible_table(pool, profile_name, &table_name).await?;
let pg_ids = candidate_rows.iter().map(|item| item.0).collect::<Vec<_>>();
let revisions_by_id = candidate_rows.into_iter().collect::<HashMap<_, _>>();
let sql = format!(
"SELECT id, to_jsonb(t) AS data FROM {} t WHERE deleted = FALSE AND id = ANY($1)",
qualified_table
@@ -1274,6 +1300,12 @@ async fn run_search(
for row in rows {
let id: i64 = row.try_get("id").unwrap_or_default();
let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
let row_revision = json_data
.get("row_revision")
.and_then(|value| value.as_i64());
if revisions_by_id.get(&id).copied() != row_revision {
continue;
}
let version = json_data.get("version").and_then(|value| value.as_i64()).unwrap_or(0);
let json_data = remap_json_to_display_names(
json_data,
@@ -1371,7 +1403,7 @@ async fn run_search(
}
}
Ok(candidates
let hits = candidates
.into_iter()
.filter_map(|candidate| {
content_map
@@ -1388,7 +1420,32 @@ async fn run_search(
archived: *archived,
})
})
.collect())
.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(
@@ -1414,10 +1471,11 @@ async fn fetch_ordered_candidate_rows(
SELECT positioned.id, to_jsonb(positioned) - 'picker_position' AS data, \
picker_position, candidate_score \
FROM positioned \
JOIN UNNEST($1::BIGINT[], $2::REAL[], $3::INTEGER[]) \
AS candidate(candidate_id, candidate_score, candidate_group) \
JOIN UNNEST($1::BIGINT[], $2::BIGINT[], $3::REAL[], $4::INTEGER[]) \
AS candidate(candidate_id, candidate_revision, candidate_score, candidate_group) \
ON candidate_id = positioned.id \
ORDER BY {} LIMIT $4 OFFSET $5",
AND candidate_revision = positioned.row_revision \
ORDER BY {} LIMIT $5 OFFSET $6",
qualified_table,
ranked_order_clause(&resolved_order, order.direction),
);
@@ -1429,6 +1487,10 @@ async fn fetch_ordered_candidate_rows(
.iter()
.map(|candidate| candidate.score)
.collect::<Vec<_>>();
let revisions = candidates
.iter()
.map(|candidate| candidate.row_revision)
.collect::<Vec<_>>();
let best_score = scores
.iter()
.copied()
@@ -1440,6 +1502,7 @@ async fn fetch_ordered_candidate_rows(
.collect::<Vec<_>>();
let rows = sqlx::query(AssertSqlSafe(sql))
.bind(&ids)
.bind(&revisions)
.bind(&scores)
.bind(&groups)
.bind(limit as i64)
@@ -1510,6 +1573,8 @@ impl Searcher for SearcherService {
mod tests {
use super::*;
use common::proto::komp_ac::search::{SearchOrder, SearchOrderDirection};
use common::search::create_search_schema;
use tantivy::query::AllQuery;
#[test]
fn search_response_mapping_exposes_aliases_only() {
@@ -1524,6 +1589,24 @@ mod tests {
assert_eq!(mapped, serde_json::json!({"customer": "Acme", "id": 4}));
}
#[test]
fn count_collector_returns_zero_documents_for_an_empty_index() {
let schema = create_search_schema();
let index = Index::create_in_ram(schema.clone());
let profile = ProfileIndex {
reader: index.reader().expect("empty reader should open"),
fields: SchemaFields::from(&schema).expect("search schema should match"),
index,
};
assert_eq!(
top_documents_for_count(&profile.reader.searcher(), &AllQuery)
.expect("empty count should not construct a zero-limit collector")
.len(),
0,
);
}
#[test]
fn search_response_mapping_fails_instead_of_leaking_a_real_name() {
let error = remap_json_to_display_names(

32
search/src/query_builder.rs
View File

@@ -1,11 +1,11 @@
use common::search::{
ARCHIVED_ROW_KEY_PREFIX, SchemaFields, json_path_term, normalize_column_name,
normalize_exact, tokenize_ngram, tokenize_word,
SchemaFields, json_path_term, normalize_column_name, normalize_exact, tokenize_ngram,
tokenize_word,
};
use common::proto::komp_ac::search::SearchVersionScope;
use tantivy::query::{
BooleanQuery, BoostQuery, EmptyQuery, FuzzyTermQuery, Occur, PhraseQuery, Query, QueryParser,
RegexQuery, TermQuery,
TermQuery,
};
use tantivy::schema::{IndexRecordOption, Term};
use tantivy::Index;
@@ -62,14 +62,15 @@ pub fn build_master_query(
));
}
let archived_rows = RegexQuery::from_pattern(
&format!("{ARCHIVED_ROW_KEY_PREFIX}.*"),
fields.row_key,
)
.map_err(|error| Status::internal(format!("Archived-row query build failed: {error}")))?;
let archived_rows = |archived| {
Box::new(TermQuery::new(
Term::from_field_u64(fields.is_archived, u64::from(archived)),
IndexRecordOption::Basic,
)) as Box<dyn Query>
};
match version_scope {
SearchVersionScope::Current => clauses.push((Occur::MustNot, Box::new(archived_rows))),
SearchVersionScope::Archived => clauses.push((Occur::Must, Box::new(archived_rows))),
SearchVersionScope::Current => clauses.push((Occur::Must, archived_rows(false))),
SearchVersionScope::Archived => clauses.push((Occur::Must, archived_rows(true))),
SearchVersionScope::All => {}
}
@@ -328,6 +329,17 @@ mod tests {
for source in documents {
let mut document = TantivyDocument::default();
document.add_u64(fields.pg_id, source.row_id);
document.add_u64(fields.row_revision, 1);
let archive = common::search::parse_archived_search_row_key(&source.row_key);
document.add_u64(fields.is_archived, u64::from(archive.is_some()));
document.add_u64(
fields.table_definition_id,
archive.map(|item| item.0 as u64).unwrap_or_default(),
);
document.add_u64(
fields.version_number,
archive.map(|item| item.2 as u64).unwrap_or_default(),
);
document.add_text(fields.table_name, source.table_name);
document.add_text(fields.row_key, &source.row_key);
let mut object = std::collections::BTreeMap::new();