Fix k-NN tests

Signed-off-by: Thomas Farr <tsfarr@amazon.com>
opensearch-project · Feb 27, 2025 · fe82e28 · fe82e28
1 parent 065fe26
commit fe82e28
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -114,6 +114,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Fixed GeoDistanceSort schema ([#805](https://github.com/opensearch-project/opensearch-api-specification/pull/805))
 - Fixed Aggregation schemas ([#801](https://github.com/opensearch-project/opensearch-api-specification/pull/801))
 - Fixed FilterQueryRequestProcessor to use correct query type ([#821](https://github.com/opensearch-project/opensearch-api-specification/pull/821))
+- Fixed `knn.train_model`'s request body `method` field to accept an object rather than a string ([#814](https://github.com/opensearch-project/opensearch-api-specification/pull/814))
 
 ### Changed
 - Changed `tasks._common:TaskInfo` and `tasks._common:TaskGroup` to be composed of a `tasks._common:TaskInfoBase` ([#683](https://github.com/opensearch-project/opensearch-api-specification/pull/683))

diff --git a/spec/namespaces/knn.yaml b/spec/namespaces/knn.yaml
@@ -390,7 +390,7 @@ components:
       in: query
       description: The default operator for query string query (AND or OR).
       schema:
-        $ref: '../schemas/knn._common.yaml#/components/schemas/DefaultOperator'
+        $ref: '../schemas/_common.query_dsl.yaml#/components/schemas/Operator'
     knn.search_models::query.df:
       name: df
       in: query
@@ -518,7 +518,7 @@ components:
       in: query
       description: Search operation type.
       schema:
-        $ref: '../schemas/knn._common.yaml#/components/schemas/SearchType'
+        $ref: '../schemas/_common.yaml#/components/schemas/SearchType'
     knn.search_models::query.seq_no_primary_term:
       name: seq_no_primary_term
       in: query
@@ -580,7 +580,7 @@ components:
       in: query
       description: Specify suggest mode.
       schema:
-        $ref: '../schemas/knn._common.yaml#/components/schemas/SuggestMode'
+        $ref: '../schemas/_common.yaml#/components/schemas/SuggestMode'
     knn.search_models::query.suggest_size:
       name: suggest_size
       in: query

diff --git a/spec/schemas/knn._common.yaml b/spec/schemas/knn._common.yaml
@@ -6,25 +6,6 @@ info:
 paths: {}
 components:
   schemas:
-    DefaultOperator:
-      type: string
-      description: The default operator for query string query (AND or OR).
-      enum:
-        - AND
-        - OR
-    SearchType:
-      type: string
-      description: Search operation type.
-      enum:
-        - dfs_query_then_fetch
-        - query_then_fetch
-    SuggestMode:
-      type: string
-      description: Specify suggest mode.
-      enum:
-        - always
-        - missing
-        - popular
     Stats:
       type: object
       properties:
@@ -143,6 +124,20 @@ components:
           $ref: '_common.yaml#/components/schemas/ByteCount'
         current_size_in_bytes:
           $ref: '_common.yaml#/components/schemas/ByteCount'
+    KnnMethod:
+      type: object
+      properties:
+        name:
+          type: string
+        space_type:
+          type: string
+        engine:
+          type: string
+        parameters:
+          type: object
+          additionalProperties: true
+      required:
+        - name
     TrainedModel:
       type: object
       properties:
@@ -166,7 +161,7 @@ components:
         compression_level:
           type: string
         method:
-          type: string
+          $ref: '#/components/schemas/KnnMethod'
         spaceType:
           type: string
       required:

diff --git a/tests/default/knn/models/search.yaml b/tests/default/knn/models/search.yaml
@@ -51,6 +51,18 @@ prologues:
               parameters:
                 code_size: 1
     status: [200]
+  - method: GET
+    path: /_plugins/_knn/models/{model_id}
+    parameters:
+      model_id: ${train_model.test_model_id}
+    retry:
+      count: 3
+    response:
+      status: 200
+      payload:
+        model_id: ${train_model.test_model_id}
+        state: created
+    status: [200]
 epilogues:
   - path: /movies
     method: DELETE

diff --git a/tests/default/knn/models/train.yaml b/tests/default/knn/models/train.yaml
@@ -70,6 +70,15 @@ chapters:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 9
+            encoder:
+              name: pq
+              parameters:
+                code_size: 3
     response:
       status: 200
     output:
@@ -104,5 +113,14 @@ chapters:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 9
+            encoder:
+              name: pq
+              parameters:
+                code_size: 3
     response:
       status: 200