Merge pull request #1 from rcarver/knn

Knn

Author: rcarver

Sources/ElasticsearchQueryBuilder/Components.swift

@@ -13,7 +13,7 @@ public protocol ArrayComponent {
 }
 
 extension ArrayComponent {
-    func makeCompactArray() -> [QueryDict] {
+    public func makeCompactArray() -> [QueryDict] {
         var value = self.makeArray()
         value.removeAll(where: \.isEmpty)
         return value
@@ -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

Sources/ElasticsearchQueryBuilder/Encoding.swift

@@ -6,6 +6,8 @@ extension QueryValue: Encodable {
         switch self {
         case let .array(value):
             try container.encode(value)
+        case let .bool(value):
+            try container.encode(value)
         case let .date(value, format: format):
             switch format {
             case .secondsSince1970:

Sources/ElasticsearchQueryBuilder/QueryValue.swift

@@ -4,6 +4,7 @@ public typealias QueryDict = [ String : QueryValue ]
 
 public enum QueryValue: Equatable {
     case array([QueryValue])
+    case bool(Bool)
     case date(Date, format: QueryDateFormat)
     case dict(QueryDict)
     case float(Float)
@@ -45,6 +46,12 @@ extension QueryValue: ExpressibleByArrayLiteral {
     }
 }
 
+extension QueryValue: ExpressibleByBooleanLiteral {
+    public init(booleanLiteral value: Bool) {
+        self = .bool(value)
+    }
+}
+
 extension QueryValue: ExpressibleByDictionaryLiteral {
     public init(dictionaryLiteral elements: (String, QueryValue)...) {
         var dict = QueryDict()

Tests/ElasticsearchQueryBuilderTests/ComponentTests.swift

@@ -175,6 +175,51 @@ 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,
+                    "index": true
+                ]
+            } 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,
+                "index": true,
+                "filter": [
+                    [
+                        "match_bool_prefix": [
+                            "message": "quick brown f"
+                        ]
+                    ]
+                ]
+            ]
+        ])
+    }
+}
+
 final class FunctionScoreTests: XCTestCase {
     func testBuild() throws {
         @ElasticsearchQueryBuilder func build() -> some esb.QueryDSL {