From e741df9c7c0e4be77d6d6887bdf8cb2ecf55ac28 Mon Sep 17 00:00:00 2001
From: Peter-Josef Meisch <pj.meisch@sothawo.com>
Date: Tue, 14 Feb 2023 20:45:16 +0100
Subject: [PATCH] Support KNN queries in NativeQuery.

Original Pull Request #2458
Closes #2433
---
 .../elasticsearch/client/elc/NativeQuery.java     | 11 +++++++++++
 .../client/elc/NativeQueryBuilder.java            | 15 +++++++++++++++
 .../client/elc/RequestConverter.java              |  6 ++++--
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
index 0dec41fb1..b4237e378 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
@@ -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;
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
index 86de08fcd..7ffe553cb 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
@@ -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);
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
index 93c29f1f3..2bd8c5334 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
@@ -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());