Add filtered knn search

elastic4s-core/src/test/scala/com/sksamuel/elastic4s/requests/searches/queries/KnnBuilderFnTest.scala

@@ -0,0 +1,21 @@
+package com.sksamuel.elastic4s.requests.searches.queries
+
+import com.sksamuel.elastic4s.handlers.searches.knn.KnnBuilderFn
+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
+
+class KnnBuilderFnTest extends AnyFunSuite with Matchers {
+
+  test("Basic Knn request generates proper query.") {
+    val request = Knn("image-vector", 50, Seq(54,10,-2))
+    KnnBuilderFn(request).string shouldBe
+      """{"field":"image-vector","query_vector":[54.0,10.0,-2.0],"k":1,"num_candidates":50,"boost":1.0}"""
+  }
+  test("Knn with all fields generates proper query.") {
+    val request = Knn("image-vector", 50, Seq(54,10,-2)) k 5 filter TermQuery("file-type", "png") similarity 10 boost .4
+    KnnBuilderFn(request).string shouldBe
+    """{"field":"image-vector","query_vector":[54.0,10.0,-2.0],"k":5,"num_candidates":50,"similarity":10.0,"filter":{"term":{"file-type":{"value":"png"}}},"boost":0.4}"""
+  }
+}

elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/requests/searches/knn/Knn.scala

@@ -1,10 +1,22 @@
 package com.sksamuel.elastic4s.requests.searches.knn
 
+import com.sksamuel.elastic4s.ext.OptionImplicits.RichOptionImplicits
+import com.sksamuel.elastic4s.requests.searches.queries.Query
+
 case class Knn(
      field: String,
      numCandidates: Int,
      queryVector: Seq[Double],
      k: Int = 1,
      similarity: Option[Float] = None,
-     boost: Double = 1.0
-)
+     filter: Option[Query] = None,
+     boost: Double = 1.0) {
+
+  def k(k: Int): Knn = copy(k = k)
+
+  def similarity(similarity: Float): Knn = copy(similarity = similarity.some)
+
+  def filter(filter: Query): Knn = copy(filter = filter.some)
+
+  def boost(boost: Double): Knn = copy(boost = boost)
+}

elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/searches/knn/KnnBuilderFn.scala

@@ -1,5 +1,6 @@
 package com.sksamuel.elastic4s.handlers.searches.knn
 
+import com.sksamuel.elastic4s.handlers.searches.queries.QueryBuilderFn
 import com.sksamuel.elastic4s.json.{XContentBuilder, XContentFactory}
 import com.sksamuel.elastic4s.requests.searches.knn.Knn
 
@@ -14,6 +15,7 @@ object KnnBuilderFn {
       case Some(value) => builder.field("similarity", value)
       case _ =>
     }
+    knn.filter.foreach(filter => builder.rawField("filter", QueryBuilderFn(filter)))
     builder.field("boost", knn.boost)
     builder.endObject()
     builder