Add support to run Knn on multiple fields (#2945)

Philippus · Nov 3, 2023 · 3b5919a · 3b5919a
1 parent bff9d65
commit 3b5919a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/...4s-core/src/main/scala/com/sksamuel/elastic4s/requests/searches/SearchBodyBuilderFn.scala b/...4s-core/src/main/scala/com/sksamuel/elastic4s/requests/searches/SearchBodyBuilderFn.scala
@@ -29,6 +29,13 @@ object SearchBodyBuilderFn {
     request.collapse.map(CollapseBuilderFn.apply).foreach(x => builder.rawField("collapse", x.string))
     request.knn.map(KnnBuilderFn.apply).foreach(x => builder.rawField("knn", x.string))
 
+    if (request.multipleKnn.nonEmpty) {
+      builder.startArray("knn")
+      val arrayBody: String = request.multipleKnn.map(KnnBuilderFn.apply).map(_.string).mkString(",")
+      builder.rawValue(arrayBody)
+      builder.endArray()
+    }
+
     request.from.foreach(builder.field("from", _))
     request.size.foreach(builder.field("size", _))
 

diff --git a/...test/scala/com/sksamuel/elastic4s/requests/searches/queries/SearchBodyBuilderFnTest.scala b/...test/scala/com/sksamuel/elastic4s/requests/searches/queries/SearchBodyBuilderFnTest.scala
@@ -6,6 +6,8 @@ import com.sksamuel.elastic4s.requests.searches.queries.geo.GeoDistanceQuery
 import com.sksamuel.elastic4s.requests.searches.sort.{GeoDistanceSort, SortOrder}
 import com.sksamuel.elastic4s.requests.searches.{GeoPoint, SearchBodyBuilderFn}
 import com.sksamuel.elastic4s.ext.OptionImplicits._
+import com.sksamuel.elastic4s.requests.searches.knn.Knn
+import com.sksamuel.elastic4s.requests.searches.term.TermQuery
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatest.matchers.should.Matchers
 
@@ -53,4 +55,15 @@ class SearchBodyBuilderFnTest extends AnyFunSuite with Matchers {
     SearchBodyBuilderFn(req).string shouldBe
       """{"query":{"geo_distance":{"distance":"100km","location":[-79.38871,43.65435]}},"size":100,"sort":[{"_geo_distance":{"location":[[-79.38871,43.65435]],"order":"asc","unit":"km"}}]}"""
   }
+
+  test("multiple Knn queries") {
+    val multipleKnnDefinition =Seq(
+      Knn("image-vector", 50, Seq(54.1, 10.2, -2.3)),
+      Knn("image-vector", 50, Seq(54.1, 10.2, -2.3)) k 5 filter TermQuery("file-type", "png") similarity 10 boost .4)
+
+    val req = search("example") multipleKnn multipleKnnDefinition
+
+    SearchBodyBuilderFn(req).string shouldBe
+      """{"knn":[{"field":"image-vector","query_vector":[54.1,10.2,-2.3],"k":1,"num_candidates":50,"boost":1.0},{"field":"image-vector","query_vector":[54.1,10.2,-2.3],"k":5,"num_candidates":50,"similarity":10.0,"filter":{"term":{"file-type":{"value":"png"}}},"boost":0.4}]}"""
+  }
 }
diff --git a/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/requests/searches/SearchRequest.scala b/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/requests/searches/SearchRequest.scala
@@ -59,7 +59,8 @@ case class SearchRequest(indexes: Indexes,
                          typedKeys: Option[Boolean] = None,
                          runtimeMappings: Seq[RuntimeMapping] = Nil,
                          ext: Map[String, Any] = Map.empty,
-                         knn: Option[Knn] = None) {
+                         knn: Option[Knn] = None,
+                         multipleKnn: Seq[Knn] = Nil) {
 
   /** Adds a single string query to this search
     *
@@ -292,4 +293,6 @@ case class SearchRequest(indexes: Indexes,
   def ext(ext: Map[String, Any]): SearchRequest = copy(ext = ext)
 
   def knn(knn: Knn): SearchRequest = copy(knn = knn.some)
+
+  def multipleKnn(multipleKnn: Iterable[Knn]): SearchRequest = copy(multipleKnn = multipleKnn.toSeq)
 }