Skip to content

Support KNN queries in NativeQuery. #2458

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 14, 2023
Merged

Support KNN queries in NativeQuery. #2458

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
*/
package org.springframework.data.elasticsearch.client.elc;

import co.elastic.clients.elasticsearch._types.KnnQuery;
import co.elastic.clients.elasticsearch._types.SortOptions;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
Expand Down Expand Up @@ -53,6 +54,7 @@ public class NativeQuery extends BaseQuery {
private List<SortOptions> sortOptions = Collections.emptyList();

private Map<String, JsonData> searchExtensions = Collections.emptyMap();
@Nullable private KnnQuery knnQuery;

public NativeQuery(NativeQueryBuilder builder) {
super(builder);
Expand All @@ -70,6 +72,7 @@ public NativeQuery(NativeQueryBuilder builder) {
"Cannot add an NativeQuery in a NativeQuery");
}
this.springDataQuery = builder.getSpringDataQuery();
this.knnQuery = builder.getKnnQuery();
}

public NativeQuery(@Nullable Query query) {
Expand Down Expand Up @@ -124,6 +127,14 @@ public void setSpringDataQuery(@Nullable org.springframework.data.elasticsearch.
this.springDataQuery = springDataQuery;
}

/**
* @since 5.1
*/
@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}

@Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery;
Expand Down
15 changes: 15 additions & 0 deletions src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
*/
package org.springframework.data.elasticsearch.client.elc;

import co.elastic.clients.elasticsearch._types.KnnQuery;
import co.elastic.clients.elasticsearch._types.SortOptions;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
Expand Down Expand Up @@ -52,6 +53,7 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
private Map<String, JsonData> searchExtensions = new LinkedHashMap<>();

@Nullable private org.springframework.data.elasticsearch.core.query.Query springDataQuery;
@Nullable private KnnQuery knnQuery;

public NativeQueryBuilder() {}

Expand Down Expand Up @@ -91,6 +93,11 @@ public Map<String, JsonData> getSearchExtensions() {
return this.searchExtensions;
}

@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}

@Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery;
Expand Down Expand Up @@ -207,6 +214,14 @@ public NativeQueryBuilder withQuery(org.springframework.data.elasticsearch.core.
return this;
}

/**
* @since 5.1
*/
public NativeQueryBuilder withKnnQuery(KnnQuery knnQuery) {
this.knnQuery = knnQuery;
return this;
}

public NativeQuery build() {
Assert.isTrue(query == null || springDataQuery == null, "Cannot have both a native query and a Spring Data query");
return new NativeQuery(this);
Expand Down
6 changes: 4 additions & 2 deletions src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1413,7 +1413,8 @@ private void prepareNativeSearch(NativeQuery query, SearchRequest.Builder builde
builder //
.suggest(query.getSuggester()) //
.collapse(query.getFieldCollapse()) //
.sort(query.getSortOptions());
.sort(query.getSortOptions()) //
.knn(query.getKnnQuery());

if (!isEmpty(query.getAggregations())) {
builder.aggregations(query.getAggregations());
Expand All @@ -1433,7 +1434,8 @@ private void prepareNativeSearch(NativeQuery query, MultisearchBody.Builder buil
builder //
.suggest(query.getSuggester()) //
.collapse(query.getFieldCollapse()) //
.sort(query.getSortOptions());
.sort(query.getSortOptions()) //
.knn(query.getKnnQuery());

if (!isEmpty(query.getAggregations())) {
builder.aggregations(query.getAggregations());
Expand Down