add KNearestNeighbor

Author: rcarver

Sources/ElasticsearchQueryBuilder/Components.swift

@@ -30,6 +30,10 @@ public struct RootComponent<Component: DictComponent>: RootQueryable, DictCompon
     }
 }
 
+public struct EmptyArrayComponent: ArrayComponent {
+    public func makeArray() -> [QueryDict] { [] }
+}
+
 /// Namespace for `@ElasticsearchQueryBuilder` components
 public enum esb {}
 
@@ -184,6 +188,45 @@ extension esb {
         }
     }
 
+    /// Adds `knn` block to the query syntax.
+    public struct KNearestNeighbor<Component: ArrayComponent>: DictComponent {
+        let field: String
+        let vector: [Double]
+        let options: QueryDict
+        var filter: Component
+        public init(
+            _ field: String,
+            _ vector: [Double],
+            options: () -> QueryDict = { [:] },
+            @QueryArrayBuilder filter: () -> Component
+        ) {
+            self.field = field
+            self.vector = vector
+            self.options = options()
+            self.filter = filter()
+        }
+        public init(
+            _ field: String,
+            _ vector: [Double],
+            options: () -> QueryDict = { [:] }
+        ) where Component == EmptyArrayComponent {
+            self.field = field
+            self.vector = vector
+            self.options = options()
+            self.filter = EmptyArrayComponent()
+        }
+        public func makeDict() -> QueryDict {
+            var dict: QueryDict = self.options
+            dict["field"] = .string(self.field)
+            dict["query_vector"] = .array(self.vector)
+            let filterValues = self.filter.makeCompactArray()
+            if !filterValues.isEmpty {
+                dict["filter"] = .array(filterValues.map(QueryValue.dict))
+            }
+            return [ "knn" : .dict(dict) ]
+        }
+    }
+
     /// Adds `function_score` block to the query syntax.
     public struct FunctionScore<Component: DictComponent>: DictComponent {
         var component: Component

Tests/ElasticsearchQueryBuilderTests/ComponentTests.swift

@@ -175,6 +175,49 @@ final class BoolTests: XCTestCase {
     }
 }
 
+final class KNearestNeighborTests: XCTestCase {
+    func testBuildBasic() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.KNearestNeighbor("vector_field", [1,2,3])
+        }
+        XCTAssertNoDifference(build().makeQuery(), [
+            "knn": [
+                "field": "vector_field",
+                "query_vector": [1.0, 2.0, 3.0],
+            ]
+        ])
+    }
+    func testBuildWithOptionsAndFilter() throws {
+        @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {
+            esb.KNearestNeighbor("vector_field", [1,2,3]) {
+                [
+                    "k": 5
+                ]
+            } filter: {
+                esb.Key("match_bool_prefix") {
+                    [
+                        "message": "quick brown f"
+                    ]
+                }
+            }
+        }
+        XCTAssertNoDifference(build().makeQuery(), [
+            "knn": [
+                "field": "vector_field",
+                "query_vector": [1.0, 2.0, 3.0],
+                "k": 5,
+                "filter": [
+                    [
+                        "match_bool_prefix": [
+                            "message": "quick brown f"
+                        ]
+                    ]
+                ]
+            ]
+        ])
+    }
+}
+
 final class FunctionScoreTests: XCTestCase {
     func testBuild() throws {
         @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {