Making k and num_candidates optional for knn search #101209

pmpailis · 2023-10-23T12:18:11Z

The aim of this PR is to make num_candidates and k optional for knn search, and provide instead some sensible defaults.

As also defined in the docs, k defines how many results we will eventually report back from each shard, while num_candidates is an exploration parameter and defines the size of the nearest neighbors queue we will compute on each shard.

k is used in two places:

On the collection phase of each shard (DfsPhase#executeKnnVectorQuery) when creating the TopScoreDocCollector for executing the knn query
On the coordinator node, when merging the knn results from different shards (through the SearchPhaseController#mergeKnnResults) in order to keep the top k.

num_cands is used solely within the KnnSearchBuilder when creating the appropriate KnnVectorQueryBuilder object.

The parameter k can default to size (or DEFAULT_SIZE if not present) - while for num_cands things are a bit more interesting :) There are mainly 2 approaches that we could follow:

Define it based on k (or size), finding a sweet-spot between recall / latency, ensuring that we will try to perform as little operations as possible
Define it based on the graph size, i.e. for smaller graphs have fewer numCands (always>= k though) - while for bigger ones we might want to allow for a few more steps to avoid local minima and try to find the true nearest neighbors. This could potentially give us more freedom in controlling (and reducing) the number of nodes we visit for each shard.

In this PR, as a first step, we will look into the 1st approach though, and make numCands depend on the incoming search request, instead of the indexed data. Based on the discussion & benchmarking results in the related issue, we'll proceed to define it as num_candidates=Math.min(1.5 * k, 10_000) (where 10_000 is the limit for num_candidates)

As there are some restrictions on the relationship between numCands and k, if we are to use any default values, in this PR we change the k and numCands variables from final and apply the same validation checks that are defined in the constructor to make sure that we either have a properly defined KnnSearchBuilder instance, or fail with an IllegalArgumentException. This is to also cover the scenario where a user might have provided numCands and size values, where numCands < size, so we want to raise an exception early on.

In the future, should we decide to make numCands defined based on the graph size, we can move the implementation & assignment over to KnnVectorQueryBuilder#doToQuery where we have info on the specific shard/graph size.

Closes #97533

…o feature/97533

github-actions · 2023-10-23T12:18:28Z

Documentation preview:

✨ Changed pages

pmpailis · 2023-10-23T12:20:19Z

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

+        int adjustedK = k != null ? k : size;
+        int adjustedNumCands = numCands != null
+            ? numCands
+            : Math.min(Math.max(NUM_CANDS_DEFAULT, (int) NUM_CANDS_MULTIPLICATIVE_FACTOR * adjustedK), NUM_CANDS_LIMIT);


This is still TBD

benwtrent · 2023-10-26T19:51:30Z

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

+        List<DfsKnnResults> mergedResults = new ArrayList<>(source.knnSearch().size());
+        for (int i = 0; i < source.knnSearch().size(); i++) {
+            Integer topK = source.knnSearch().get(i).k();
+            int size = source.size() == -1 ? DEFAULT_SIZE : source.size();
+            TopDocs mergedTopDocs = TopDocs.merge(topK != null ? topK : size, topDocsLists.get(i).toArray(new TopDocs[0]));


I wonder if we could handle this within the SearchSourceBuilder directly? When parsing (thus creating via an API), parsing the knn section or size section could set the k for the knn there.

Similarly if somebody calls SearchSourceBuilder#setSize(int) that could propagate down to the knn object if it exists?

I was looking into it initially, but I wasn't impressed with a couple of approaches I've tried, as we were introducing dependencies between different component in the SearchSourceBuilder - especially when there's no guaranty of order. On the other hand, as things are now, if someone calls SearchSourceBuilder#setSize at a later point, then they would also have the responsibility of setting up & updating k which is not right 🤔 (similarly for org.elasticsearch.search.DefaultSearchContext#size(int) )

Will look into it a bit more.

Updated the functionality to go through the SearchSourceBuilder in 9092f37. Haven't set it in the SearchSourceBuilder#size method, but rather when we parse the initial size param (iff present).

Happy to update it as well, but was wondering which cases related to knn search we'd capture (was looking at the usages but I'm not very familiar with all cases).

benwtrent · 2023-10-26T19:53:51Z

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

+        this.k = in.readOptionalVInt();
+        this.numCands = in.readOptionalVInt();


when you are ready for further review, etc. all these will have to be backwards compatible. Meaning looking at the in.getTransportVersion() and verifying its at least the version where this change is introduced.

Similar logic down in writeTo(StreamOutput) method.

Thanks for the heads up @benwtrent ! Will update accordingly.

elasticsearchmachine · 2023-11-08T15:46:59Z

Hi @pmpailis, I've created a changelog YAML for you.

…o feature/97533

elasticsearchmachine · 2024-01-11T13:00:03Z

Pinging @elastic/es-search (Team:Search)

benwtrent

Great stuff!

docs/reference/rest-api/common-parms.asciidoc

...c/yamlRestTest/resources/rest-api-spec/test/search.vectors/150_knn_search_missing_params.yml

benwtrent · 2024-01-12T11:55:06Z

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -193,6 +193,7 @@ static TransportVersion def(int id) {
    public static final TransportVersion ENRICH_ELASTICSEARCH_VERSION_REMOVED = def(8_560_00_0);
    public static final TransportVersion NODE_STATS_REQUEST_SIMPLIFIED = def(8_561_00_0);

+    public static final TransportVersion KNN_K_NUMCANDS_AS_OPTIONAL_PARAMS = def(8_993_00_0);


Yes, please update it.

benwtrent · 2024-01-12T12:03:01Z

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchRequestParser.java

@@ -194,6 +196,7 @@ public void toSearchRequest(SearchRequestBuilder builder) {

    // visible for testing
    static class KnnSearch {


We should not update this class anylonger. Its deprecated and folks shouldn't use the _knn_search API.

Rolled back all changes to KnnSearch in 1aa708b

server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java

benwtrent · 2024-01-12T12:10:47Z

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

+    private void adjustKnnSizeParam() {
+        if (size > 0) knnSearch.forEach(knn -> knn.requestSize(size));
+    }


It seems like this isn't being used at all within the KnnSearchBuilder. However, we can do better.

We should have a KnnSearchBuilder.Builder that KnnSearchBuilder.fromXContent returns. That can have setters for k that gets adjusted after the last field is read when parsing the SearchSourceBuilder.

This way we have a KnnSearchBuilder.Builder#build(int) where int is the request size (the default or user provided). And KnnSearchBuilder will never have null k nor null num_candidates.

I see what you mean - Nice suggestion!
Currently, org.elasticsearch.search.builder.SearchSourceBuilder#parseXContent doesn't do much outside of the while loop other than iterating through all objects. I guess we should add the call to the builder within this parsing section though, as only then the SearchSourceBuilder would have been finalized parsing.

Do you think that we should hide all other public test usages/constructors of KnnSearchBuilder behind the builder (except from the one accepting a StreamInput), or just isolate this very specific use-case?

Do you think that we should hide all other public test usages/constructors of KnnSearchBuilder behind the builder (except from the one accepting a StreamInput), or just isolate this very specific use-case?

I think the ctor should be package private & we use a builder object for all final object values.

And KnnSearchBuilder will never have null k nor null num_candidates.

So, since the responsibility to fill out missing values will be moved over to the Builder on the node that received the original request, we wouldn't need any transport versions or backwards compatibility as well, right? The values will have already been filled when sending over the objects, so the behavior right after builder#build will be pretty much the same.

If the KnnSearchBuilder doesn't have null values, it won't need bwc serialization as nothing changes there. But, the KnnVectorQueryBuilder probably still needs it as num_candidates can be null when executed as knn-as-a-query

Yeap, knn-query would still need it. But we can take it out of knn-search. Getting on it : )

hey @pmpailis are you blocked by me here?

Nope @benwtrent ! Was just taking a look at toXContent to understand under which circumstances this is invoked, and more specifically in the context of knn-search. Will have an update for that shortly :)

…ionContext

benwtrent

Just a couple of things. Looking REALLY nice :)

benwtrent · 2024-01-30T19:27:48Z

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

+        public KnnSearchBuilder build(int size) {
+            int requestSize = size < 0 ? DEFAULT_SIZE : size;
+            int adjustedK = k == null ? requestSize : k;
+            int adjustedNumCandidates = numCandidates == null
+                ? Math.round(Math.min(NUM_CANDS_LIMIT, NUM_CANDS_MULTIPLICATIVE_FACTOR * adjustedK))
+                : numCandidates;
+            return new KnnSearchBuilder(


benwtrent · 2024-01-30T19:27:56Z

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

-    final int k;
-    final int numCands;
-    final Float similarity;
+    int k;
+    int numCands;
+    Float similarity;


This can all still be final right?

Yeap - restored them in 390aff2

test/framework/src/main/java/org/elasticsearch/test/AbstractQueryVectorBuilderTestCase.java

…eryVectorBuilderTestCase.java Co-authored-by: Benjamin Trent <ben.w.trent@gmail.com>

server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java

…o feature/97533

pmpailis added 9 commits October 9, 2023 19:22

updating

7eba0ce

temp commit

b5b2b90

Merge branch 'elastic:main' into feature/97533

0860626

Merge branch 'feature/97533' of github.com:pmpailis/elasticsearch int…

c31b2b7

…o feature/97533

updating

f9c1345

updating docs

d4e8445

updating

7cd412d

updating

20554ac

updating

343acbc

elasticsearchmachine added the v8.12.0 label Oct 23, 2023

pmpailis commented Oct 23, 2023

View reviewed changes

pmpailis added 5 commits October 23, 2023 19:07

updating

eaade90

Merge branch 'elastic:main' into feature/97533

4b5ce54

updating

a0f8152

updating

ec07bfc

updating

d24ad28

benwtrent reviewed Oct 26, 2023

View reviewed changes

pmpailis added 5 commits November 2, 2023 20:25

Merge remote-tracking branch 'origin/main' into feature/97533

2586f9c

adding requestSize to SearchExecutionContext

bf7ef86

merging main

a66f965

updating

642f481

updating

c3794fc

benwtrent added >enhancement :Search/Vectors Vector search labels Nov 8, 2023

pmpailis added 3 commits November 8, 2023 17:46

Update docs/changelog/101209.yaml

7e29689

updating

8cadde7

Merge branch 'feature/97533' of github.com:pmpailis/elasticsearch int…

c26da26

…o feature/97533

benwtrent self-requested a review January 11, 2024 15:59

pmpailis changed the title ~~[WIP] Making k and num_candidates optional for knn search~~ Making k and num_candidates optional for knn search Jan 12, 2024

benwtrent reviewed Jan 12, 2024

View reviewed changes

pmpailis added 10 commits January 12, 2024 15:27

pr iter - updating documentation

8729217

pr iter - removing unecessary check from tests

a1237a5

pr iter - rolling back changes to deprecated class

1aa708b

pr iter - setting default requestSize to DEFAULT_SIZE in SearchExecut…

4082cf7

…ionContext

minor iter

78a7fc0

Merge remote-tracking branch 'origin/main' into feature/97533

2d760f7

minor iter

d1b457c

Merge remote-tracking branch 'origin/main' into feature/97533

2b5d15d

iter

6b1c94e

iter

cfee774

benwtrent reviewed Jan 30, 2024

View reviewed changes

pmpailis and others added 5 commits January 30, 2024 23:30

Update test/framework/src/main/java/org/elasticsearch/test/AbstractQu…

d1a072d

…eryVectorBuilderTestCase.java Co-authored-by: Benjamin Trent <ben.w.trent@gmail.com>

Update test/framework/src/main/java/org/elasticsearch/test/AbstractQu…

7fc1157

…eryVectorBuilderTestCase.java Co-authored-by: Benjamin Trent <ben.w.trent@gmail.com>

Merge remote-tracking branch 'origin/main' into feature/97533

208a483

restoring final variables & updating TransportVersion

390aff2

renaming TransportVersion

8965d58

benwtrent approved these changes Jan 31, 2024

View reviewed changes

server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java Outdated Show resolved Hide resolved

pmpailis added 5 commits February 1, 2024 12:29

Merge branch 'main' into feature/97533

69f92f6

updating exception message

d2cb4b2

Merge branch 'feature/97533' of github.com:pmpailis/elasticsearch int…

9121b82

…o feature/97533

checkstyle

f964fa8

Update KnnVectorQueryBuilder.java

5d195e3

pmpailis merged commit 7ce8d76 into elastic:main Feb 1, 2024
14 checks passed

jedrazb pushed a commit to jedrazb/elasticsearch that referenced this pull request Feb 2, 2024

Making k and num_candidates optional for knn search (elastic#101209)

0df8308

mayya-sharipova mentioned this pull request Feb 15, 2024

Add retrievers using the parser-only approach #105470

Merged

pquentin mentioned this pull request Apr 8, 2024

Make k and num_candidates optional in kNN queries elastic/elasticsearch-specification#2490

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Making k and num_candidates optional for knn search #101209

Making k and num_candidates optional for knn search #101209

pmpailis commented Oct 23, 2023 •

edited

github-actions bot commented Oct 23, 2023

pmpailis Oct 23, 2023

benwtrent Oct 26, 2023

pmpailis Oct 26, 2023

pmpailis Jan 11, 2024

benwtrent Oct 26, 2023

pmpailis Oct 26, 2023

elasticsearchmachine commented Nov 8, 2023

elasticsearchmachine commented Jan 11, 2024

benwtrent left a comment

benwtrent Jan 12, 2024

benwtrent Jan 12, 2024

pmpailis Jan 12, 2024

benwtrent Jan 12, 2024

pmpailis Jan 12, 2024

benwtrent Jan 12, 2024

pmpailis Jan 12, 2024

benwtrent Jan 12, 2024

pmpailis Jan 12, 2024

benwtrent Jan 17, 2024

pmpailis Jan 17, 2024

benwtrent left a comment

benwtrent Jan 30, 2024

benwtrent Jan 30, 2024

pmpailis Jan 30, 2024

		this.k = in.readOptionalVInt();
		this.numCands = in.readOptionalVInt();

		@@ -194,6 +196,7 @@ public void toSearchRequest(SearchRequestBuilder builder) {

		// visible for testing
		static class KnnSearch {

Making k and num_candidates optional for knn search #101209

Making k and num_candidates optional for knn search #101209

Conversation

pmpailis commented Oct 23, 2023 • edited

github-actions bot commented Oct 23, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elasticsearchmachine commented Nov 8, 2023

elasticsearchmachine commented Jan 11, 2024

benwtrent left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benwtrent left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pmpailis commented Oct 23, 2023 •

edited