fix: make kNN num_candidates configurable and improve default recall

- Add knnNumCandidatesMultiplier to NaturalLanguageSearchConfiguration
  schema (default 2). Operators can tune recall vs latency by adjusting
  this value in openmetadata.yaml under naturalLanguageSearch.
- Change buildNativeESQuery to accept a numCandidatesMultiplier param;
  add no-arg overload that uses DEFAULT_KNN_NUM_CANDIDATES_MULTIPLIER (2)
- ElasticSearchVectorService.search() reads the multiplier from config
  via resolveNumCandidatesMultiplier(), falling back to the default
- Remove nestedTags path from appendFilterMustClauses: entity indices use
  flat tags objects, so nested tags query is never correct here
- Remove appendNested/appendOneNestedQuery dead code
- Update tests to match new num_candidates = max(k * multiplier, 100)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: joaopamaral

openmetadata-service/src/main/java/org/openmetadata/service/search/vector/ElasticSearchVectorService.java

@@ -18,7 +18,10 @@
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.schema.EntityInterface;
+import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.events.lifecycle.EntityLifecycleEventDispatcher;
+import org.openmetadata.service.search.SearchRepository;
 import org.openmetadata.service.search.vector.client.EmbeddingClient;
 import org.openmetadata.service.search.vector.utils.DTOs.VectorSearchResponse;
 
@@ -102,11 +105,12 @@ public VectorSearchResponse search(
         overFetchSize = Math.min(overFetchSize, k);
       }
 
+      int numCandidatesMultiplier = resolveNumCandidatesMultiplier();
       String indexName = getIndexAlias();
       while (!exhausted && byParent.size() < requestedParents) {
         String queryJson =
             VectorSearchQueryBuilder.buildNativeESQuery(
-                queryVector, overFetchSize, rawOffset, k, filters);
+                queryVector, overFetchSize, rawOffset, k, filters, numCandidatesMultiplier);
         String responseBody = executeGenericRequest("POST", "/" + indexName + "/_search", queryJson);
 
         JsonNode root = MAPPER.readTree(responseBody);
@@ -300,6 +304,18 @@ public void partialUpdateEntity(
     }
   }
 
+  private static int resolveNumCandidatesMultiplier() {
+    SearchRepository repo = Entity.getSearchRepository();
+    if (repo == null) {
+      return VectorSearchQueryBuilder.DEFAULT_KNN_NUM_CANDIDATES_MULTIPLIER;
+    }
+    NaturalLanguageSearchConfiguration cfg = repo.getSearchConfiguration().getNaturalLanguageSearch();
+    if (cfg == null || cfg.getKnnNumCandidatesMultiplier() == null) {
+      return VectorSearchQueryBuilder.DEFAULT_KNN_NUM_CANDIDATES_MULTIPLIER;
+    }
+    return cfg.getKnnNumCandidatesMultiplier();
+  }
+
   public void close() {
     try {
       if (client != null && client._transport() != null) {

openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorSearchQueryBuilder.java

@@ -14,6 +14,7 @@ public class VectorSearchQueryBuilder {
   private static final Logger LOG = LoggerFactory.getLogger(VectorSearchQueryBuilder.class);
   private static final String ANY = "__ANY__";
   private static final String NONE = "__NONE__";
+  static final int DEFAULT_KNN_NUM_CANDIDATES_MULTIPLIER = 2;
 
   /** Build a full search request body (size + _source + query) for standalone vector search. */
   public static String build(
@@ -74,7 +75,17 @@ private static void appendKnnQuery(
 
   public static String buildNativeESQuery(
       float[] vector, int size, int from, int k, Map<String, List<String>> filters) {
-    int numCandidates = Math.max(k, 100);
+    return buildNativeESQuery(vector, size, from, k, filters, DEFAULT_KNN_NUM_CANDIDATES_MULTIPLIER);
+  }
+
+  public static String buildNativeESQuery(
+      float[] vector,
+      int size,
+      int from,
+      int k,
+      Map<String, List<String>> filters,
+      int numCandidatesMultiplier) {
+    int numCandidates = Math.max(k * numCandidatesMultiplier, 100);
     StringBuilder sb =
         new StringBuilder(512)
             .append("{\"size\":")
@@ -92,24 +103,15 @@ public static String buildNativeESQuery(
             .append(numCandidates);
 
     sb.append(",\"filter\":{\"bool\":{\"must\":[");
-    appendFilterMustClauses(sb, filters, true);
+    appendFilterMustClauses(sb, filters);
     sb.append("]}}"); // close must array and bool
 
     sb.append("}}"); // close knn object
     return sb.toString();
   }
 
-  private static void appendFilterMustClauses(StringBuilder sb, Map<String, List<String>> filters) {
-    appendFilterMustClauses(sb, filters, false);
-  }
-
-  /**
-   * @param nestedTags when true, emits a nested query for "tags" (required by the ES-native vector
-   *     index where tags is mapped as nested); when false, emits a flat terms query (used by
-   *     OpenSearch, which queries the regular entity indices where tags is a plain object).
-   */
   private static void appendFilterMustClauses(
-      StringBuilder sb, Map<String, List<String>> filters, boolean nestedTags) {
+      StringBuilder sb, Map<String, List<String>> filters) {
     sb.append("{\"term\":{\"deleted\":false}}");
     for (var e : filters.entrySet()) {
       String field = e.getKey();
@@ -126,11 +128,7 @@ private static void appendFilterMustClauses(
           }
           case "tags" -> {
             sb.append(',');
-            if (nestedTags) {
-              appendNested(sb, "tags", "tags.tagFQN", values);
-            } else {
-              appendFlat(sb, "tags.tagFQN", values);
-            }
+            appendFlat(sb, "tags.tagFQN", values);
           }
           case "domains" -> {
             sb.append(',');
@@ -170,35 +168,6 @@ private static void appendFilterMustClauses(
     }
   }
 
-  private static void appendNested(StringBuilder sb, String path, String field, List<String> vals) {
-    sb.append("{\"nested\":{\"path\":\"").append(path).append("\",\"query\":");
-    if (vals.size() == 1) {
-      appendOneNestedQuery(sb, field, vals.get(0));
-    } else {
-      sb.append("{\"bool\":{\"should\":[");
-      for (int i = 0; i < vals.size(); i++) {
-        if (i > 0) sb.append(',');
-        appendOneNestedQuery(sb, field, vals.get(i));
-      }
-      sb.append("]}}");
-    }
-    sb.append("}}");
-  }
-
-  private static void appendOneNestedQuery(StringBuilder sb, String field, String val) {
-    switch (val) {
-      case ANY -> sb.append("{\"exists\":{\"field\":\"").append(field).append("\"}}");
-      case NONE -> sb.append("{\"bool\":{\"must_not\":{\"exists\":{\"field\":\"")
-          .append(field)
-          .append("\"}}}}");
-      default -> sb.append("{\"term\":{\"")
-          .append(field)
-          .append("\":\"")
-          .append(escape(val))
-          .append("\"}}");
-    }
-  }
-
   private static void appendFlat(StringBuilder sb, String field, List<String> vals) {
     if (vals.size() == 1) {
       appendOneFlat(sb, field, vals.get(0));

openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorSearchQueryBuilderTest.java

@@ -736,15 +736,20 @@ void testNativeESQueryTopLevelKnnStructure() throws Exception {
   void testNativeESQueryNumCandidates() throws Exception {
     float[] vector = {0.1f};
 
-    // k < 100 → num_candidates should be 100
-    String query1 = VectorSearchQueryBuilder.buildNativeESQuery(vector, 10, 0, 50, Map.of());
+    // default multiplier (2): k * 2 < 100 → num_candidates should be 100
+    String query1 = VectorSearchQueryBuilder.buildNativeESQuery(vector, 10, 0, 30, Map.of());
     JsonNode root1 = MAPPER.readTree(query1);
     assertEquals(100, root1.get("knn").get("num_candidates").asInt());
 
-    // k > 100 → num_candidates should equal k
+    // default multiplier (2): k * 2 > 100 → num_candidates should be k * 2
     String query2 = VectorSearchQueryBuilder.buildNativeESQuery(vector, 10, 0, 200, Map.of());
     JsonNode root2 = MAPPER.readTree(query2);
-    assertEquals(200, root2.get("knn").get("num_candidates").asInt());
+    assertEquals(400, root2.get("knn").get("num_candidates").asInt());
+
+    // custom multiplier (5): num_candidates = max(k * 5, 100)
+    String query3 = VectorSearchQueryBuilder.buildNativeESQuery(vector, 10, 0, 100, Map.of(), 5);
+    JsonNode root3 = MAPPER.readTree(query3);
+    assertEquals(500, root3.get("knn").get("num_candidates").asInt());
   }
 
   @Test
@@ -815,8 +820,8 @@ void testNativeESQueryWithTagsFilter() throws Exception {
 
     assertEquals(2, mustFilters.size());
     JsonNode tagsFilter = mustFilters.get(1);
-    assertTrue(tagsFilter.has("nested"));
-    assertEquals("tags", tagsFilter.get("nested").get("path").asText());
+    assertTrue(tagsFilter.has("term"));
+    assertEquals("PII.Sensitive", tagsFilter.get("term").get("tags.tagFQN").asText());
   }
 
   @Test

openmetadata-spec/src/main/resources/json/schema/configuration/elasticSearchConfiguration.json

@@ -159,6 +159,12 @@
           "default": 10,
           "minimum": 1
         },
+        "knnNumCandidatesMultiplier": {
+          "description": "Multiplier applied to k when computing num_candidates for Elasticsearch kNN vector search. num_candidates = max(k * multiplier, 100). Higher values improve recall at the cost of latency. Defaults to 2.",
+          "type": "integer",
+          "default": 2,
+          "minimum": 1
+        },
         "providerClass": {
           "description": "Fully qualified class name of the NLQService implementation to use",
           "type": "string",