diff --git a/apps/hash-graph/src/subcommand/server.rs b/apps/hash-graph/src/subcommand/server.rs index 938060aad69..e693ea4f1f8 100644 --- a/apps/hash-graph/src/subcommand/server.rs +++ b/apps/hash-graph/src/subcommand/server.rs @@ -25,6 +25,7 @@ use hash_graph_authorization::policies::store::{PolicyStore, PrincipalStore}; use hash_graph_embeddings::{OpenAiEmbeddingClient, OpenAiEmbeddingClientConfig}; use hash_graph_postgres_store::store::{ DatabaseConnectionInfo, DatabasePoolConfig, PostgresStorePool, PostgresStoreSettings, + SemanticSearchSettings, }; use hash_graph_store::{filter::protection::PropertyProtectionFilterConfig, pool::StorePool}; use hash_graph_type_fetcher::FetchingPool; @@ -227,6 +228,28 @@ pub struct ServerConfig { #[clap(long, env = "HASH_GRAPH_SKIP_EMBEDDING_CREATION")] pub skip_embedding_creation: bool, + /// Candidates a semantic search ranks per requested result. + /// + /// Raising this recovers neighbours the quantized ranking misordered, at the cost of reading + /// more full vectors when re-scoring. + #[clap( + long, + default_value_t = SemanticSearchSettings::default().candidate_overfetch, + env = "HASH_GRAPH_SEMANTIC_SEARCH_CANDIDATE_OVERFETCH", + )] + pub semantic_search_candidate_overfetch: NonZero, + + /// Lower bound for the size of a semantic search's vector-index walk. + /// + /// Searches asking for few results otherwise walk the index too shallowly to find the + /// neighbours their re-scoring could order. + #[clap( + long, + default_value_t = SemanticSearchSettings::default().minimum_ef_search, + env = "HASH_GRAPH_SEMANTIC_SEARCH_MINIMUM_EF_SEARCH", + )] + pub semantic_search_minimum_ef_search: usize, + /// Disables filter protection that prevents enumeration attacks on protected properties. /// /// When enabled (protection disabled), queries filtering on protected properties like email @@ -524,6 +547,10 @@ pub async fn server(mut args: ServerArgs) -> Result<(), Report> { } else { PropertyProtectionFilterConfig::hash_default() }, + semantic_search: SemanticSearchSettings { + candidate_overfetch: args.config.semantic_search_candidate_overfetch, + minimum_ef_search: args.config.semantic_search_minimum_ef_search, + }, }, ) .await diff --git a/libs/@local/graph/postgres-store/src/store/mod.rs b/libs/@local/graph/postgres-store/src/store/mod.rs index b90ee497b15..2ed318f3f4f 100644 --- a/libs/@local/graph/postgres-store/src/store/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/mod.rs @@ -10,7 +10,8 @@ pub use self::{ postgres::{ AsClient, BeginReadOnlyTransaction, Context, InTransaction, IsolationLevel, NoTransaction, PostgresStore, PostgresStorePool, PostgresStoreSettings, PostgresStoreTransactionBuilder, - Transaction, TransactionBuilder, TransactionOptions, TransactionState, + SemanticSearchSettings, Transaction, TransactionBuilder, TransactionOptions, + TransactionState, }, validation::{StoreCache, StoreProvider}, }; diff --git a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/search.rs b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/search.rs index 5353ecac7d6..c51050ee616 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/search.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/search.rs @@ -24,10 +24,7 @@ use type_system::{ use crate::store::{ AsClient, PostgresStore, - postgres::{ - InTransaction, - query::{QUANTIZED_RANK_OVERFETCH, SelectCompiler}, - }, + postgres::{InTransaction, query::SelectCompiler}, }; const fn empty_response() -> SearchEntitiesResponse { @@ -159,7 +156,7 @@ where let temporal_axes = QueryTemporalAxesUnresolved::live_only().resolve(); let request_filter = request_filter(&entity_type_ids, &web_ids); - let candidate_pool = limit.saturating_mul(QUANTIZED_RANK_OVERFETCH); + let candidate_pool = self.settings.semantic_search.candidate_pool(limit); self.prepare_hnsw_scan(candidate_pool).await?; // The branches may permit overlapping sets of entities, so the candidate keys @@ -218,13 +215,13 @@ where /// /// An HNSW scan stops after `ef_search` tuples (default 40) regardless of the statement's /// limit. The iterative mode resumes the walk when the filters discard candidates, or when - /// the pool exceeds the setting's range of 1 to 1000. `SET LOCAL` scopes the settings to - /// the enclosing transaction. + /// the pool exceeds the setting's range. `SET LOCAL` scopes the settings to the enclosing + /// transaction. async fn prepare_hnsw_scan(&self, candidate_pool: usize) -> Result<(), Report> { let settings = format!( "SET LOCAL hnsw.ef_search = {}; SET LOCAL hnsw.iterative_scan = relaxed_order;", - candidate_pool.clamp(1, 1000) + self.settings.semantic_search.ef_search(candidate_pool) ); self.as_client() .batch_execute(&settings) diff --git a/libs/@local/graph/postgres-store/src/store/postgres/mod.rs b/libs/@local/graph/postgres-store/src/store/postgres/mod.rs index cb58ef2d4f7..9de8ab189ce 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/mod.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/mod.rs @@ -8,7 +8,9 @@ mod seed_policies; mod traversal_context; use alloc::{borrow::Cow, sync::Arc}; -use core::{borrow::Borrow, fmt::Debug, hash::Hash, marker::PhantomData}; +use core::{ + borrow::Borrow, fmt::Debug, hash::Hash, marker::PhantomData, num::NonZero, ops::RangeInclusive, +}; use std::collections::{HashMap, HashSet}; use error_stack::{Report, ResultExt as _, TryReportStreamExt as _}; @@ -88,6 +90,55 @@ use crate::store::error::{ StoreError, VersionedUrlAlreadyExists, }; +/// The range `hnsw.ef_search` accepts. +const HNSW_EF_SEARCH_RANGE: RangeInclusive = 1..=1000; + +/// How a semantic search walks the vector index before its results are re-scored exactly. +#[derive(Debug, Clone, Copy)] +pub struct SemanticSearchSettings { + /// Multiplier from a search's limit to the number of candidates ranked on the quantized + /// embedding. + /// + /// The exact re-scoring only reorders the candidates the walk produced, so this decides + /// whether a true neighbour can reach the result at all. It covers the quantization's + /// misordering alone — filters participate in the ranking itself, never in the re-scoring. + pub candidate_overfetch: NonZero, + + /// Lower bound for `hnsw.ef_search`. + /// + /// Deriving the search list from the candidate pool alone leaves the walk shallow for + /// searches asking for few results, which is where a missed neighbour is most likely. + pub minimum_ef_search: usize, +} + +impl Default for SemanticSearchSettings { + fn default() -> Self { + Self { + candidate_overfetch: NonZero::new(4).expect("4 is not zero"), + minimum_ef_search: 400, + } + } +} + +impl SemanticSearchSettings { + /// The number of candidates to rank for a search returning at most `limit` records. + #[must_use] + pub const fn candidate_pool(&self, limit: usize) -> usize { + limit.saturating_mul(self.candidate_overfetch.get()) + } + + /// The size of the search list a walk over `candidate_pool` candidates uses. + /// + /// A scan stops once its list is exhausted, so the list must hold the whole pool; the + /// iterative scan resumes the walk beyond the setting's upper bound. + #[must_use] + pub fn ef_search(&self, candidate_pool: usize) -> usize { + candidate_pool + .max(self.minimum_ef_search) + .clamp(*HNSW_EF_SEARCH_RANGE.start(), *HNSW_EF_SEARCH_RANGE.end()) + } +} + #[derive(Debug)] pub struct PostgresStoreSettings { pub validate_links: bool, @@ -97,6 +148,7 @@ pub struct PostgresStoreSettings { /// When set, filters on protected properties will automatically exclude /// specified entity types to prevent enumeration attacks. pub filter_protection: PropertyProtectionFilterConfig<'static>, + pub semantic_search: SemanticSearchSettings, } impl Default for PostgresStoreSettings { @@ -105,6 +157,7 @@ impl Default for PostgresStoreSettings { validate_links: true, skip_embedding_creation: false, filter_protection: PropertyProtectionFilterConfig::hash_default(), + semantic_search: SemanticSearchSettings::default(), } } } @@ -4385,3 +4438,74 @@ where Ok(()) } } + +#[cfg(test)] +mod tests { + use core::num::NonZero; + + use super::{HNSW_EF_SEARCH_RANGE, SemanticSearchSettings}; + + fn settings(candidate_overfetch: usize, minimum_ef_search: usize) -> SemanticSearchSettings { + SemanticSearchSettings { + candidate_overfetch: NonZero::new(candidate_overfetch) + .expect("the test overfetch should be non-zero"), + minimum_ef_search, + } + } + + #[test] + fn candidate_pool_scales_with_the_limit() { + let settings = settings(4, 400); + + assert_eq!(settings.candidate_pool(0), 0); + assert_eq!(settings.candidate_pool(1), 4); + assert_eq!(settings.candidate_pool(100), 400); + assert_eq!( + settings.candidate_pool(usize::MAX), + usize::MAX, + "an unreachable limit should saturate instead of wrapping" + ); + } + + #[test] + fn ef_search_holds_the_whole_pool() { + let settings = settings(4, 1); + + assert_eq!(settings.ef_search(40), 40); + assert_eq!( + settings.ef_search(4000), + *HNSW_EF_SEARCH_RANGE.end(), + "a pool beyond the setting's range should cap, leaving the rest to the iterative scan" + ); + } + + #[test] + fn ef_search_keeps_the_walk_deep_for_small_pools() { + let settings = settings(4, 400); + + assert_eq!( + settings.ef_search(4), + 400, + "a search for a single record should still walk deeply" + ); + assert_eq!( + settings.ef_search(800), + 800, + "a pool above the floor should decide the walk itself" + ); + } + + #[test] + fn ef_search_stays_within_the_accepted_range() { + assert_eq!( + settings(4, 0).ef_search(0), + *HNSW_EF_SEARCH_RANGE.start(), + "an empty pool should still emit a value the setting accepts" + ); + assert_eq!( + settings(4, usize::MAX).ef_search(0), + *HNSW_EF_SEARCH_RANGE.end(), + "an unreachable floor should cap instead of being emitted" + ); + } +} diff --git a/libs/@local/graph/postgres-store/src/store/postgres/ontology/entity_type.rs b/libs/@local/graph/postgres-store/src/store/postgres/ontology/entity_type.rs index 002570ac82e..1efdd1251bc 100644 --- a/libs/@local/graph/postgres-store/src/store/postgres/ontology/entity_type.rs +++ b/libs/@local/graph/postgres-store/src/store/postgres/ontology/entity_type.rs @@ -73,8 +73,7 @@ use crate::store::{ crud::{QueryIndices, QueryRecordDecode, TypedRow}, ontology::{PostgresOntologyOwnership, read::OntologyTypeTraversalData}, query::{ - Distinctness, PostgresRecord, PostgresSorting, QUANTIZED_RANK_OVERFETCH, - ReferenceTable, SelectCompiler, Table, + Distinctness, PostgresRecord, PostgresSorting, ReferenceTable, SelectCompiler, Table, }, }, validation::StoreProvider, @@ -1247,7 +1246,7 @@ where compiler .rank_by_quantized_distance(&embedding_path, &embedding) .change_context(QueryError)?; - let candidate_pool = limit.saturating_mul(QUANTIZED_RANK_OVERFETCH); + let candidate_pool = self.settings.semantic_search.candidate_pool(limit); compiler.set_limit(candidate_pool); // Same shape as the entity search: the candidate CTE ranks on the quantized embedding, diff --git a/tests/graph/benches/graph/scenario/runner.rs b/tests/graph/benches/graph/scenario/runner.rs index 27987ec59c5..ae2cce22da5 100644 --- a/tests/graph/benches/graph/scenario/runner.rs +++ b/tests/graph/benches/graph/scenario/runner.rs @@ -217,6 +217,7 @@ impl Runner { validate_links: true, skip_embedding_creation: true, filter_protection: PropertyProtectionFilterConfig::new(), // Disabled for benchmarks + ..PostgresStoreSettings::default() }, ) .await