This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 750551f3 [integrations][vector-store][java] Apply Elasticsearch KNN
filters before selecting neighbors (#1001)
750551f3 is described below
commit 750551f3c5a13aa8c67cfe5b6f173b6f46ee0e60
Author: Weiqing Yang <[email protected]>
AuthorDate: Thu Aug 13 21:09:10 2026 -0700
[integrations][vector-store][java] Apply Elasticsearch KNN filters before
selecting neighbors (#1001)
Generated-by: Claude Code 2.1.228 (Claude Opus 5)
---
docs/content/docs/development/vector_stores.md | 2 +-
.../elasticsearch/ElasticsearchVectorStore.java | 40 ++++++++++++++-------
.../ElasticsearchVectorStoreTest.java | 42 ++++++++++++++++++++++
3 files changed, 70 insertions(+), 14 deletions(-)
diff --git a/docs/content/docs/development/vector_stores.md
b/docs/content/docs/development/vector_stores.md
index 07d5eafa..f4462a07 100644
--- a/docs/content/docs/development/vector_stores.md
+++ b/docs/content/docs/development/vector_stores.md
@@ -800,7 +800,7 @@ Elasticsearch is currently supported in the Java API only.
To use Elasticsearch
| `dims` | int | `768` | Vector dimensionality
|
| `k` | int | None | Number of nearest
neighbors to return; can be overridden per query |
| `num_candidates` | int | None | Candidate set size
for ANN search; can be overridden per query |
-| `filter_query` | str | None | Raw JSON
Elasticsearch filter query (DSL) applied as a post-filter |
+| `filter_query` | str | None | Raw JSON
Elasticsearch filter query (DSL) restricting which documents a KNN query can
match |
| `host` | str | `"http://localhost:9200"` | Elasticsearch
endpoint |
| `hosts` | str | None | Comma-separated list
of Elasticsearch endpoints |
| `username` | str | None | Username for basic
authentication |
diff --git
a/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
b/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
index 424c9d39..5d66e479 100644
---
a/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
+++
b/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
@@ -82,8 +82,8 @@ import java.util.*;
* <li>{@code k} (optional): Number of nearest neighbors to return; can be
overridden per query.
* <li>{@code num_candidates} (optional): Candidate set size for ANN search;
can be overridden per
* query.
- * <li>{@code filter_query} (optional): A raw JSON Elasticsearch filter
query (DSL) that is
- * applied as a post-filter; can be overridden per query.
+ * <li>{@code filter_query} (optional): A raw JSON Elasticsearch filter
query (DSL) restricting
+ * which documents a KNN query can match; can be overridden per query.
* <li>{@code host} or {@code hosts} (optional): Elasticsearch endpoint(s).
If omitted, defaults
* to {@code localhost:9200}.
* <li>Authentication (optional): Either basic auth via {@code
username}/{@code password}, or API
@@ -572,7 +572,9 @@ public class ElasticsearchVectorStore extends
BaseVectorStore
*
* <p>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 class ElasticsearchVectorStore extends
BaseVectorStore
List<Float> 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<Map<String, Object>> 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 e89365e1..2b281483 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 class ElasticsearchVectorStoreTest {
((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<Document> 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<Document> 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.
