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
@@ -516,6 +539,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/migrations/graph-migrations/v008__entity_types/up.sql

@@ -36,5 +36,8 @@ CREATE TABLE entity_type_constrains_link_destinations_on (
 CREATE TABLE entity_type_embeddings (
     ontology_id UUID PRIMARY KEY REFERENCES entity_types,
     embedding VECTOR(3072) NOT NULL,
-    updated_at_transaction_time TIMESTAMP WITH TIME ZONE NOT NULL
+    updated_at_transaction_time TIMESTAMP WITH TIME ZONE NOT NULL,
+    embedding_bits BIT(3072) NOT NULL GENERATED ALWAYS AS (
+        binary_quantize(embedding)::BIT(3072)
+    ) STORED
 );

libs/@local/graph/migrations/graph-migrations/v009__entities/up.sql

@@ -156,8 +156,15 @@ CREATE TABLE entity_embeddings (
     embedding VECTOR(3072) NOT NULL,
     updated_at_decision_time TIMESTAMP WITH TIME ZONE NOT NULL,
     updated_at_transaction_time TIMESTAMP WITH TIME ZONE NOT NULL,
+    embedding_bits BIT(3072) NOT NULL GENERATED ALWAYS AS (
+        binary_quantize(embedding)::BIT(3072)
+    ) STORED,
     FOREIGN KEY (web_id, entity_uuid) REFERENCES entity_ids
 );
 
 CREATE UNIQUE INDEX entity_embeddings_idx
 ON entity_embeddings (web_id, entity_uuid, property) NULLS NOT DISTINCT;
+
+-- Only the combined per-entity embedding: the per-property rows are a different space.
+CREATE INDEX entity_embeddings_hnsw
+ON entity_embeddings USING hnsw (embedding_bits bit_hamming_ops) WHERE property IS NULL;

libs/@local/graph/postgres-store/postgres_migrations/V57__entity_embedding_bits.sql

@@ -0,0 +1,12 @@
+ALTER TABLE entity_embeddings
+    ADD COLUMN embedding_bits bit(3072) NOT NULL
+    GENERATED ALWAYS AS (binary_quantize(embedding)::bit(3072)) STORED;
+
+ALTER TABLE entity_type_embeddings
+    ADD COLUMN embedding_bits bit(3072) NOT NULL
+    GENERATED ALWAYS AS (binary_quantize(embedding)::bit(3072)) STORED;
+
+-- Only the combined per-entity embedding: the per-property rows are a different space.
+CREATE INDEX entity_embeddings_hnsw ON entity_embeddings
+    USING hnsw (embedding_bits bit_hamming_ops)
+    WHERE property IS NULL;

libs/@local/graph/postgres-store/src/snapshot/entity/batch.rs

@@ -81,9 +81,12 @@ where
                         (LIKE entity_edge INCLUDING ALL)
                         ON COMMIT DROP;
 
+                    -- Staging holds the rows for a single copy, so it needs neither the HNSW
+                    -- index nor the quantized column derived from the vectors.
                     CREATE TEMPORARY TABLE entity_embeddings_tmp
-                        (LIKE entity_embeddings INCLUDING ALL)
+                        (LIKE entity_embeddings INCLUDING ALL EXCLUDING INDEXES)
                         ON COMMIT DROP;
+                    ALTER TABLE entity_embeddings_tmp DROP COLUMN embedding_bits;
                 ",
             )
             .instrument(tracing::info_span!(
@@ -235,8 +238,26 @@ where
                     INSERT INTO entity_edge
                         SELECT * FROM entity_edge_tmp;
 
-                    INSERT INTO entity_embeddings
-                        SELECT * FROM entity_embeddings_tmp;
+                    -- The explicit list leaves out the generated `embedding_bits` column, which
+                    -- rejects inserts even when the staging table is empty.
+                    INSERT INTO entity_embeddings (
+                        web_id,
+                        entity_uuid,
+                        draft_id,
+                        property,
+                        embedding,
+                        updated_at_decision_time,
+                        updated_at_transaction_time
+                    )
+                        SELECT
+                            web_id,
+                            entity_uuid,
+                            draft_id,
+                            property,
+                            embedding,
+                            updated_at_decision_time,
+                            updated_at_transaction_time
+                        FROM entity_embeddings_tmp;
             ",
             )
             .instrument(tracing::info_span!(

libs/@local/graph/postgres-store/src/snapshot/ontology/entity_type/batch.rs

@@ -38,6 +38,9 @@ where
                     CREATE TEMPORARY TABLE entity_type_embeddings_tmp (
                         LIKE entity_type_embeddings INCLUDING ALL
                     ) ON COMMIT DROP;
+                    -- Staging holds the rows for a single copy, so it does not need the
+                    -- quantized column derived from the vectors.
+                    ALTER TABLE entity_type_embeddings_tmp DROP COLUMN embedding_bits;
                 ",
             )
             .instrument(tracing::info_span!(
@@ -102,8 +105,18 @@ where
                     INSERT INTO entity_types
                         SELECT * FROM entity_types_tmp;
 
-                    INSERT INTO entity_type_embeddings
-                        SELECT * FROM entity_type_embeddings_tmp;
+                    -- The explicit list leaves out the generated `embedding_bits` column, which
+                    -- rejects inserts even when the staging table is empty.
+                    INSERT INTO entity_type_embeddings (
+                        ontology_id,
+                        embedding,
+                        updated_at_transaction_time
+                    )
+                        SELECT
+                            ontology_id,
+                            embedding,
+                            updated_at_transaction_time
+                        FROM entity_type_embeddings_tmp;
                 ",
             )
             .instrument(tracing::info_span!(

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/mod.rs

@@ -2,6 +2,7 @@ mod delete;
 pub(crate) mod provenance;
 mod query;
 mod read;
+mod search;
 mod summary;
 mod table;
 
@@ -11,7 +12,6 @@ use std::collections::{HashMap, HashSet};
 
 use error_stack::{FutureExt as _, Report, ResultExt as _, TryReportStreamExt as _, ensure};
 use futures::{StreamExt as _, TryStreamExt as _, stream};
-use hash_codec::numeric::Real;
 use hash_graph_authorization::policies::{
     Authorized, MergePolicies, PolicyComponents, Request, RequestContext, ResourceId,
     action::ActionName,
@@ -29,9 +29,9 @@ use hash_graph_store::{
         EntityValidationReport, EntityValidationType, HasPermissionForEntitiesParams,
         PatchEntityParams, QueryConversion, QueryEntitiesParams, QueryEntitiesResponse,
         QueryEntitiesTableParams, QueryEntitiesTableResponse, QueryEntitySubgraphParams,
-        QueryEntitySubgraphResponse, SearchEntitiesFilter, SearchEntitiesParams,
-        SearchEntitiesResponse, SummarizeEntitiesParams, SummarizeEntitiesResponse,
-        UpdateEntityEmbeddingsParams, ValidateEntityComponents, ValidateEntityParams,
+        QueryEntitySubgraphResponse, SearchEntitiesParams, SearchEntitiesResponse,
+        SummarizeEntitiesParams, SummarizeEntitiesResponse, UpdateEntityEmbeddingsParams,
+        ValidateEntityComponents, ValidateEntityParams,
     },
     entity_type::{EntityTypeStore as _, IncludeEntityTypeOption},
     error::{
@@ -630,11 +630,7 @@ where
         // always match their key row (foreign keys, and the write paths maintaining
         // `entity_edition_cache` in the same transaction) and the distinct key pins all
         // row-multiplying columns.
-        //
-        // TODO(BE-618): revisit the embedding shape once the distance query is index-backed.
-        if !compiler.has_embeddings_filter() {
-            compiler.set_statement_shape(StatementShape::KeysFirst);
-        }
+        compiler.set_statement_shape(StatementShape::KeysFirst);
 
         let cursor_parameters = params.sorting.encode().change_context(QueryError)?;
         let cursor_indices = params
@@ -1850,95 +1846,20 @@ where
         actor_id: ActorEntityUuid,
         params: SearchEntitiesParams,
     ) -> Result<SearchEntitiesResponse, Report<QueryError>> {
-        let SearchEntitiesParams {
-            embedding,
-            maximum_semantic_distance,
-            limit,
-            include_entity_types,
-            filter:
-                SearchEntitiesFilter {
-                    entity_type_ids,
-                    web_ids,
-                    include_drafts,
-                },
-        } = params;
-
-        // TODO(BE-618): optimize the query — it scans embeddings without a vector index. A
-        //   halfvec/HNSW index needs an ANN-friendly query shape to be usable (the current
-        //   `MIN(<=>) GROUP BY` defeats it). The returned entities can also be trimmed to the
-        //   fields the search bar and inference actually use, but the query is the bottleneck.
-        let maximum_distance =
-            Real::try_from(maximum_semantic_distance.into_inner()).change_context(QueryError)?;
-
-        // The search always runs against the current time and never returns archived entities.
-        let mut filters = vec![
-            Filter::CosineDistance(
-                FilterExpression::Path {
-                    path: EntityQueryPath::Embedding,
-                },
-                FilterExpression::Parameter {
-                    parameter: Parameter::Vector(embedding),
-                    convert: None,
-                },
-                FilterExpression::Parameter {
-                    parameter: Parameter::Decimal(maximum_distance),
-                    convert: None,
-                },
-            ),
-            Filter::Equal(
-                FilterExpression::Path {
-                    path: EntityQueryPath::Archived,
-                },
-                FilterExpression::Parameter {
-                    parameter: Parameter::Boolean(false),
-                    convert: None,
-                },
-            ),
-        ];
+        // The search issues several statements — one candidate read per policy branch, the
+        // rerank, and the hydration. Under `READ COMMITTED` each would use its own MVCC
+        // snapshot, so a write committing in between could rank an entity that the hydration no
+        // longer returns.
+        let transaction = self
+            .begin_read_only_transaction()
+            .await
+            .change_context(QueryError)?;
 
-        if !entity_type_ids.is_empty() {
-            filters.push(Filter::Any(
-                entity_type_ids
-                    .iter()
-                    .map(Filter::for_entity_by_type_id)
-                    .collect(),
-            ));
-        }
-        if !web_ids.is_empty() {
-            filters.push(Filter::In(
-                FilterExpression::Path {
-                    path: EntityQueryPath::WebId,
-                },
-                FilterExpressionList::ParameterList {
-                    parameters: ParameterList::WebIds(&web_ids),
-                },
-            ));
-        }
+        let response = transaction.search_entities_impl(actor_id, params).await?;
 
-        let response = self
-            .query_entities(
-                actor_id,
-                QueryEntitiesParams {
-                    filter: Filter::All(filters),
-                    temporal_axes: QueryTemporalAxesUnresolved::live_only(),
-                    sorting: EntityQuerySorting {
-                        paths: vec![],
-                        cursor: None,
-                    },
-                    conversions: Vec::new(),
-                    limit,
-                    include_drafts,
-                    include_entity_types: include_entity_types
-                        .then_some(IncludeEntityTypeOption::Closed),
-                    include_permissions: false,
-                },
-            )
-            .await?;
+        transaction.commit().await.change_context(QueryError)?;
 
-        Ok(SearchEntitiesResponse {
-            entities: response.entities,
-            closed_multi_entity_types: response.closed_multi_entity_types,
-        })
+        Ok(response)
     }
 
     #[tracing::instrument(level = "info", skip(self, params))]

libs/@local/graph/postgres-store/src/store/postgres/knowledge/entity/query.rs

@@ -145,10 +145,6 @@ impl QueryRecordDecode for Entity {
             draft_id: row.get(indices.draft_id),
         };
 
-        if let Ok(distance) = row.try_get::<_, f64>("distance") {
-            tracing::trace!(%entity_id, %distance, "Entity embedding was calculated");
-        }
-
         let property_metadata = row
             .get::<_, Option<serde_json::Value>>(indices.property_metadata)
             .map(|value| {