BE-743: Make the semantic-search candidate pool and ef_search configurable

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<usize>,
+
+    /// 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<GraphError>> {
             } 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

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},
 };

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<QueryError>> {
         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)

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<usize> = 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<usize>,
+
+    /// 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"
+        );
+    }
+}

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,

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