Add 'knn' section to search endpoint by jtibshirani · Pull Request #88002 · elastic/elasticsearch

jtibshirani · 2022-06-24T06:43:02Z

This PR adds a new knn section to the search request:

POST products/_search
{
  "query": {
    "multi_match": {
      "query": "black flower dress",
      "fields": ["title", "description"],
    },
    "boost": 0.9
  },
  "knn": {
    "field": "title_vector",
    "query_vector": [0.3f, 0.1f, ...],
    "k": 5,
    "num_candidates": 50,
    "boost": 0.1
  },
  "size": 20
}

As explained in #87625, the search returns a global top k documents across all
shards:

In the DFS phase, collect the top k vector matches per shard, combine them
and keep the global top k.
Then move to the query phase. Convert the top k vector matches into a new
KnnScoreDocQueryBuilder that matches only those top documents. For the final
query, take a boolean disjunction between this KnnScoreDocQueryBuilder query and
the search request query.

Commits:

Rename KnnSearchRequestBuilder -> KnnSearchRequestParser to avoid confusion with new KnnSearchBuilder
Add KnnScoreDocQueryBuilder
Implement kNN search using DFS phase

Addresses #87625.

jtibshirani · 2022-06-24T06:44:51Z

NOTE: this PR targets a feature branch knn-search. There are still several TODOs (tracked in #87625) like support for filters.

elasticmachine · 2022-06-24T16:29:52Z

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

elasticmachine · 2022-06-24T16:47:58Z

Pinging @elastic/clients-team (Team:Clients)

javanna · 2022-06-24T18:41:52Z

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

+    static void adjustSearchType(SearchRequest searchRequest) {
+        // if there's a kNN search, always use DFS_QUERY_THEN_FETCH
+        if (searchRequest.hasKnnSearch()) {
+            searchRequest.searchType(DFS_QUERY_THEN_FETCH);


drive-by comment: do we want to return an error in case search type is explicitly set and differs from dfs query then fetch? Maybe this is something to think about for later, does not look high priority in this PR. Or maybe you already though about it.

Thanks @javanna, I meant to add a comment about this. Yes in a follow-up PR I plan to disallow setting the search type explicitly when knn is provided too.

javanna · 2022-06-24T18:51:43Z

server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java

+            // Emulate TransportSearchAction logic: first adjust search type, then check minimize roundtrips
+            TransportSearchAction.adjustSearchType(searchRequest);
+
+            // Minimize roundtrips should always be false because kNN uses DFS_QUERY_THEN_FETCH


too bad that we can't minimize roundtrips! maybe that's a good reminder to rethink whether we can do something about this, I will try to think about this.

In a follow-up, I hope to optimize the case where there is only kNN, with no query or aggregations. Then we can run kNN during the main query phase (using QUERY_THEN_FETCH), and it will be possible to minimize roundtrips.

Also, I found this test to be a little fragile/ hacky -- I'm happy to remove it if people don't think it's valuable. I checked that the new CCS REST tests already cover the important cases.

I see, agreed that it's a bit hacky because it requires adjustSearchType to be called. I would rather add a very simple unit test around adjustSearchType then. this specific test method already tests the scenario with dfs enabled, so we should just verify that the search type is adjusted when knn is used.

Also, Nhat is running benchmarks to asses how bad it is that we don't minimize roundtrips in some cases, in the context of #73971 . I would imagine that the results of the benchmark will also inform whether we need to put more effort into minimizing roundtrips when knn is used.

👍 I'll rework this test. Looking forward to seeing Nhat's results.

sethmlarson

API looks good, one thought:

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/40_knn_search.yml

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

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

...tegTest/groovy/org/elasticsearch/gradle/fixtures/AbstractGradleInternalPluginFuncTest.groovy

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

jtibshirani · 2022-06-28T08:50:41Z

@elasticmachine run elasticsearch-ci/packaging-tests-unix-sample

jtibshirani · 2022-06-28T09:36:57Z

Thanks @mayya-sharipova for the helpful review!

@jpountz do you have time to review just to check the overall strategy looks okay? As a reminder, the PR targets a feature branch knn-search. There are still several TODOs (tracked in #87625) like support for filters.

jtibshirani · 2022-06-28T15:06:26Z

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

+            searchRequest.searchType(QUERY_THEN_FETCH);
+        }
+
+        // if there's a kNN search, always use DFS_QUERY_THEN_FETCH


In a follow-up I'm thinking of handling this in SearchSourceBuilder#rewrite instead.

jpountz

I only skimmed through the PR. One thing that looks a bit odd to me if I read the code correctly is that we compute global nearest vectors on the coordinating node and then send them all to all shards, where many of them might get ignored as part of QueryBuilder#toQuery. Would it be possible to only send the relevant vectors to each shard?

jpountz · 2022-06-28T16:15:07Z

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

        final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = resultConsumer.reduce();
-        final boolean queryAndFetchOptimization = queryResults.length() == 1;
+        final boolean queryAndFetchOptimization = queryResults.length() == 1
+            && context.getRequest().searchType() == SearchType.QUERY_THEN_FETCH;


Why would we not run the fetch phase and the query phase in the same roundtrip when the search type has a DFS phase?

Previously, when there was a single shard we would always set the search type to QUERY_THEN_FETCH:

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

Lines 988 to 992 in 8fb440d

// optimize search type for cases where there is only one shard group to search on

if (shardIterators.size() == 1) {

// if we only have one group, then we always want Q_T_F, no need for DFS, and no need to do THEN since we hit one shard

searchRequest.searchType(QUERY_THEN_FETCH);

}

This makes sense, since you don't need DFS with one shard. Now, we might still execute DFS_QUERY_THEN_FETCH with a single shard, when kNN is used. From failing tests, I saw that we don't apply the "query and fetch" optimization in the query phase when DFS is enabled. (Surprisingly, we use different codepath to execute the query phase when DFS is enabled vs. not. There are two different actions QUERY_ID_ACTION_NAME and QUERY_ACTION_NAME).

Here was my thinking for next steps:

In a follow-up, I plan to use QUERY_THEN_FETCH for kNN with a single shard. This will remove the special handling, since we will always know QUERY_THEN_FETCH is used when there's one shard.

For now I'll change this check to queryResults.length() == 1 && context.getRequest().hasKnnSearch() == false and add a comment explaining.

Thanks for explaining, the proposed follow-up sounds good to me.

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

server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

jtibshirani · 2022-06-28T17:04:15Z

One thing that looks a bit odd to me if I read the code correctly is that we compute global nearest vectors on the coordinating node and then send them all to all shards, where many of them might get ignored as part of QueryBuilder#toQuery.

That's correct! I did this because I thought it was a nice invariant that we always send the same QueryBuilder to each shard. I can't think of a prior case where we've sent a different query to each shard? It didn't seem like a big deal because k is bounded and typically not that large.

jpountz · 2022-06-28T17:19:15Z

The closest thing that I know of is how searches sorted by descending @timestamp fan out to shards that hold recent data first and then to shards that hold older data with a search_after key that hopefully sometimes helps skip these shards entirely in the common case when the most recent documents were all in the recent shards.

I agree that it shouldn't be a big deal in practice. One could argue that it is k times the number of shards, but the size of these queries is still likely smaller than the responses that shards will send back.

jtibshirani · 2022-06-28T17:30:12Z

@jpountz thanks for your comments! It sounds like things look okay overall, no major flags. In the interest of time, I'm going to continue with these PRs (merging into a feature branch) and tag both you and @mayya-sharipova for a final review before merging to main.

jtibshirani · 2022-06-29T10:01:33Z

I pushed some updates:

Add comment to clarify why the 'query and fetch' optimization doesn't work when kNN is enabled. It'd be best to support this -- I'll address this later when we tackle the TODO item "Optimize the single shard case".
Only send the relevant score docs to each shard. This can be more efficient, especially when some shards don't return any kNN results (maybe because the filter doesn't match?) It also made things cleaner since I could avoid the new methods Lucene.writeScoreDocWithShardIndex. We do something similar in SearchQueryThenFetchAsyncAction#rewriteShardSearchRequest, so I used the same method name.

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

sethmlarson

Looks good from an API perspective.

I'm guessing that the filter property of knn is being implemented in a future PR since there are no YAML tests cases with the property? If so that's fine.

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

mayya-sharipova

@jtibshirani Thanks, new changes LGTM@

This PR adds a new `knn` section to the search request. As explained in #87625, the search returns a global top k documents across all shards: * In the DFS phase, collect the top k vector matches per shard, combine them and keep the global top k. * Then move to the query phase. Convert the top k vector matches into a new `KnnScoreDocQueryBuilder` that matches only those top documents. For the final query, take a boolean disjunction between this `KnnScoreDocQueryBuilder` query and the search request `query`. Commits: * Rename KnnSearchRequestBuilder -> KnnSearchRequestParser to avoid confusion with new KnnSearchBuilder * Add `KnnScoreDocQueryBuilder` * Implement kNN search using DFS phase Addresses #87625.

jtibshirani added 2 commits June 23, 2022 23:31

Rename KnnSearchRequestBuilder -> KnnSearchRequestParser

a5744fd

Add ScoreDocQuery

ceb6eb3

jtibshirani force-pushed the knn-search-phase branch from 303726f to fa8697b Compare June 24, 2022 07:19

Implement kNN search within DFS phase.

4d78eb4

jtibshirani force-pushed the knn-search-phase branch from fa8697b to 4d78eb4 Compare June 24, 2022 07:30

jtibshirani added >feature :Search/Search Search-related issues that do not fall into other categories labels Jun 24, 2022

jtibshirani marked this pull request as ready for review June 24, 2022 16:29

elasticmachine added the Team:Search Meta label for search team label Jun 24, 2022

jtibshirani requested a review from mayya-sharipova June 24, 2022 16:30

sethmlarson added the Team:Clients Meta label for clients team label Jun 24, 2022

javanna reviewed Jun 24, 2022

View reviewed changes

jtibshirani added the v8.4.0 label Jun 24, 2022

sethmlarson reviewed Jun 24, 2022

View reviewed changes

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/40_knn_search.yml Show resolved Hide resolved

mayya-sharipova reviewed Jun 25, 2022

View reviewed changes

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

mayya-sharipova reviewed Jun 25, 2022

View reviewed changes

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

jtibshirani added >non-issue and removed >feature labels Jun 25, 2022

Fixes to KnnSearchBuilder

6e0e039

jtibshirani force-pushed the knn-search-phase branch from 149017a to 6e0e039 Compare June 25, 2022 04:07

mayya-sharipova reviewed Jun 26, 2022

View reviewed changes

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java Outdated Show resolved Hide resolved

jtibshirani added 2 commits June 27, 2022 14:34

Refactor adjustSearchType logic

4533a2a

Merge remote-tracking branch 'upstream/master' into knn-search-phase

be6d502

mayya-sharipova reviewed Jun 27, 2022

View reviewed changes

...tegTest/groovy/org/elasticsearch/gradle/fixtures/AbstractGradleInternalPluginFuncTest.groovy Show resolved Hide resolved

Fix REST BWC tests

9724355

mayya-sharipova reviewed Jun 27, 2022

View reviewed changes

server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java Show resolved Hide resolved

Address code review comments

e023d80

jtibshirani mentioned this pull request Jun 28, 2022

Integrate ANN into _search endpoint #87625

Closed

8 tasks

jtibshirani commented Jun 28, 2022

View reviewed changes

jpountz reviewed Jun 28, 2022

View reviewed changes

jtibshirani added 3 commits June 29, 2022 10:39

Only send relevant score docs to each shard

70cf63e

Add comment around 'query and fetch' optimization

95202bb

Merge branch 'knn-search' into knn-search-phase

382e527

jtibshirani requested review from mayya-sharipova and sethmlarson June 29, 2022 10:31

jpountz reviewed Jun 29, 2022

View reviewed changes

server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java Show resolved Hide resolved

Implement KnnScoreDocQueryBuilder#rewrite and clean up tests

86ea977

sethmlarson reviewed Jun 29, 2022

View reviewed changes

mayya-sharipova reviewed Jun 29, 2022

View reviewed changes

server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java Show resolved Hide resolved

mayya-sharipova approved these changes Jun 30, 2022

View reviewed changes

jtibshirani merged commit d53d01f into elastic:knn-search Jun 30, 2022

jtibshirani deleted the knn-search-phase branch June 30, 2022 14:55

sethmlarson mentioned this pull request Jun 30, 2022

Add the knn property to search requests elastic/elasticsearch-specification#1779

Closed

jtibshirani added :Search Relevance/Vectors Vector search and removed :Search/Search Search-related issues that do not fall into other categories labels Jul 21, 2022

jtibshirani mentioned this pull request Jul 21, 2022

Integrate ANN into _search endpoint #88694

Merged

sethmlarson mentioned this pull request Jul 26, 2022

Add 'knn' property to Search request body elastic/elasticsearch-specification#1792

Merged

jtibshirani mentioned this pull request Jul 28, 2022

Avoid extra roundtrips in ANN search #88921

Open

	// optimize search type for cases where there is only one shard group to search on
	if (shardIterators.size() == 1) {
	// if we only have one group, then we always want Q_T_F, no need for DFS, and no need to do THEN since we hit one shard
	searchRequest.searchType(QUERY_THEN_FETCH);
	}

Conversation

jtibshirani commented Jun 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jtibshirani commented Jun 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

elasticmachine commented Jun 24, 2022

Uh oh!

elasticmachine commented Jun 24, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jtibshirani Jun 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sethmlarson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jtibshirani commented Jun 28, 2022

Uh oh!

jtibshirani commented Jun 28, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jpountz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jtibshirani Jun 29, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jtibshirani commented Jun 28, 2022

Uh oh!

jpountz commented Jun 28, 2022

Uh oh!

jtibshirani commented Jun 28, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jtibshirani commented Jun 29, 2022

Uh oh!

Uh oh!

sethmlarson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mayya-sharipova left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

jtibshirani commented Jun 24, 2022 •

edited

Loading

jtibshirani commented Jun 24, 2022 •

edited

Loading

jtibshirani Jun 24, 2022 •

edited

Loading

jtibshirani Jun 29, 2022 •

edited

Loading

jtibshirani commented Jun 28, 2022 •

edited

Loading