Authentication (optional): Either basic auth via {@code username}/{@code password}, or API
@@ -572,7 +572,9 @@ private void deleteDocuments(
*
* The method prepares a KNN search request using the supplied {@code embedding} and merges
* default arguments from the store with the provided {@code args}. Optional filter queries
- * (JSON DSL) are applied as a post filter.
+ * (JSON DSL) restrict the documents the KNN search may match, so the nearest neighbours are
+ * selected from among the matching documents rather than filtered out afterwards. Up to {@code
+ * k} matching documents are returned even when the closest vectors overall do not match.
*
* @param embedding The embedding vector to search with
* @param limit Maximum number of items the caller is interested in; used as a fallback for
@@ -603,20 +605,32 @@ public List queryEmbedding(
List queryVector = new ArrayList<>(embedding.length);
for (float v : embedding) queryVector.add(v);
+ final String finalCombined = combined;
SearchRequest.Builder builder =
new SearchRequest.Builder()
.index(index)
.knn(
- kb ->
- kb.field(this.vectorField)
- .queryVector(queryVector)
- .k(k)
- .numCandidates(numCandidates));
-
- if (combined != null) {
- final String finalCombined = combined;
- builder = builder.postFilter(f -> f.withJson(new StringReader(finalCombined)));
- }
+ kb -> {
+ kb.field(this.vectorField)
+ .queryVector(queryVector)
+ .k(k)
+ .numCandidates(numCandidates);
+ // Filter inside the KNN clause rather than after it, so the
+ // k nearest neighbours are chosen from the documents that
+ // match. A post-filter can only discard hits the vector
+ // search already picked, which yields fewer than k results
+ // whenever the nearest vectors belong to filtered-out
+ // documents.
+ if (finalCombined != null) {
+ kb.filter(
+ f ->
+ f.withJson(
+ new StringReader(
+ finalCombined)));
+ }
+ return kb;
+ });
+
final SearchResponse> searchResponse =
(SearchResponse) this.client.search(builder.build(), Map.class);
diff --git a/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
index e89365e1f..2b281483a 100644
--- a/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
+++ b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
@@ -191,6 +191,48 @@ public void testFiltersDsl() throws Exception {
((CollectionManageableVectorStore) store).deleteCollection(name);
}
+ @Test
+ public void testQueryEmbeddingFiltersBeforeSelectingNeighbors() throws Exception {
+ // Contract: filters restrict the candidate set of the KNN search itself, so a matching
+ // document is returned even when it is not among the k nearest vectors overall. Applying
+ // the filter after the KNN phase instead would return nothing here, because the k nearest
+ // vectors all belong to the other user.
+ String name = "knn_prefilter";
+ ((CollectionManageableVectorStore) store).createCollectionIfNotExists(name, Map.of());
+
+ List docs = new ArrayList<>();
+ // Six documents pointing the same way as the query vector, none of them alice's.
+ for (int i = 0; i < 6; i++) {
+ Document bob =
+ new Document("bob document " + i, Map.of("user_id", "bob"), "doc_bob_" + i);
+ bob.setEmbedding(new float[] {1.0f, 0.0f, 0.0f, 0.0f, 0.0f});
+ docs.add(bob);
+ }
+ // Three alice documents pointing orthogonally, so they never make the unfiltered top k.
+ for (int i = 0; i < 3; i++) {
+ Document alice =
+ new Document(
+ "alice document " + i, Map.of("user_id", "alice"), "doc_alice_" + i);
+ alice.setEmbedding(new float[] {0.0f, 0.0f, 0.0f, 0.0f, 1.0f});
+ docs.add(alice);
+ }
+ store.addEmbedding(docs, name, Collections.emptyMap());
+ Thread.sleep(1000);
+
+ List alice =
+ store.queryEmbedding(
+ new float[] {1.0f, 0.0f, 0.0f, 0.0f, 0.0f},
+ 5,
+ name,
+ Map.of("user_id", "alice"),
+ Collections.emptyMap());
+
+ Assertions.assertEquals(3, alice.size());
+ Assertions.assertTrue(alice.stream().allMatch(d -> d.getId().startsWith("doc_alice_")));
+
+ ((CollectionManageableVectorStore) store).deleteCollection(name);
+ }
+
@Test
public void testUpdateOverwritesExistingDocument() throws Exception {
// ES bulk index is upsert by id — update should rewrite the doc in place.