Support KNN queries in NativeQuery.

Original Pull Request #2458
Closes #2433
This commit is contained in:
Peter-Josef Meisch 2023-02-14 20:45:16 +01:00 committed by GitHub
parent 3875bb6abf
commit e741df9c7c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 30 additions and 2 deletions

View File

@ -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;
@ -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);
@ -70,6 +72,7 @@ public class NativeQuery extends BaseQuery {
"Cannot add an NativeQuery in a NativeQuery");
}
this.springDataQuery = builder.getSpringDataQuery();
this.knnQuery = builder.getKnnQuery();
}
public NativeQuery(@Nullable Query query) {
@ -124,6 +127,14 @@ public class NativeQuery extends BaseQuery {
this.springDataQuery = springDataQuery;
}
/**
* @since 5.1
*/
@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}
@Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery;

View File

@ -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;
@ -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() {}
@ -91,6 +93,11 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
return this.searchExtensions;
}
@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}
@Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery;
@ -207,6 +214,14 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
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);

View File

@ -1413,7 +1413,8 @@ class RequestConverter {
builder //
.suggest(query.getSuggester()) //
.collapse(query.getFieldCollapse()) //
.sort(query.getSortOptions());
.sort(query.getSortOptions()) //
.knn(query.getKnnQuery());
if (!isEmpty(query.getAggregations())) {
builder.aggregations(query.getAggregations());
@ -1433,7 +1434,8 @@ class RequestConverter {
builder //
.suggest(query.getSuggester()) //
.collapse(query.getFieldCollapse()) //
.sort(query.getSortOptions());
.sort(query.getSortOptions()) //
.knn(query.getKnnQuery());
if (!isEmpty(query.getAggregations())) {
builder.aggregations(query.getAggregations());