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; 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.SortOptions;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation; import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.Query; 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 List<SortOptions> sortOptions = Collections.emptyList();
private Map<String, JsonData> searchExtensions = Collections.emptyMap(); private Map<String, JsonData> searchExtensions = Collections.emptyMap();
@Nullable private KnnQuery knnQuery;
public NativeQuery(NativeQueryBuilder builder) { public NativeQuery(NativeQueryBuilder builder) {
super(builder); super(builder);
@ -70,6 +72,7 @@ public class NativeQuery extends BaseQuery {
"Cannot add an NativeQuery in a NativeQuery"); "Cannot add an NativeQuery in a NativeQuery");
} }
this.springDataQuery = builder.getSpringDataQuery(); this.springDataQuery = builder.getSpringDataQuery();
this.knnQuery = builder.getKnnQuery();
} }
public NativeQuery(@Nullable Query query) { public NativeQuery(@Nullable Query query) {
@ -124,6 +127,14 @@ public class NativeQuery extends BaseQuery {
this.springDataQuery = springDataQuery; this.springDataQuery = springDataQuery;
} }
/**
* @since 5.1
*/
@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}
@Nullable @Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() { public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery; return springDataQuery;

View File

@ -15,6 +15,7 @@
*/ */
package org.springframework.data.elasticsearch.client.elc; 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.SortOptions;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation; import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.Query; 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<>(); private Map<String, JsonData> searchExtensions = new LinkedHashMap<>();
@Nullable private org.springframework.data.elasticsearch.core.query.Query springDataQuery; @Nullable private org.springframework.data.elasticsearch.core.query.Query springDataQuery;
@Nullable private KnnQuery knnQuery;
public NativeQueryBuilder() {} public NativeQueryBuilder() {}
@ -91,6 +93,11 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
return this.searchExtensions; return this.searchExtensions;
} }
@Nullable
public KnnQuery getKnnQuery() {
return knnQuery;
}
@Nullable @Nullable
public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() { public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
return springDataQuery; return springDataQuery;
@ -207,6 +214,14 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
return this; return this;
} }
/**
* @since 5.1
*/
public NativeQueryBuilder withKnnQuery(KnnQuery knnQuery) {
this.knnQuery = knnQuery;
return this;
}
public NativeQuery build() { public NativeQuery build() {
Assert.isTrue(query == null || springDataQuery == null, "Cannot have both a native query and a Spring Data query"); Assert.isTrue(query == null || springDataQuery == null, "Cannot have both a native query and a Spring Data query");
return new NativeQuery(this); return new NativeQuery(this);

View File

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