Fix k-NN model training tests to work with small training sample set

Signed-off-by: Thomas Farr <[email protected]>
opensearch-project · Feb 27, 2025 · 6f7acf8 · 6f7acf8
1 parent f419abc
commit 6f7acf8
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/tests/default/knn/models.yaml b/tests/default/knn/models.yaml
@@ -27,6 +27,10 @@ prologues:
         - {recommendation_vector: [1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], duration: 12.2}
         - {index: {_index: movies, _id: '2'}}
         - {recommendation_vector: [2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5], duration: 7.1}
+        - {index: {_index: movies, _id: '3'}}
+        - {recommendation_vector: [3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5], duration: 4.8}
+        - {index: {_index: movies, _id: '4'}}
+        - {recommendation_vector: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5], duration: 18.6}
     status: [200]
   - method: POST
     path: /_plugins/_knn/models/model-1/_train
@@ -41,6 +45,15 @@ prologues:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 4
+            encoder:
+              name: pq
+              parameters:
+                code_size: 2
     status: [200]
 epilogues:
   - path: /movies

diff --git a/tests/default/knn/models/search.yaml b/tests/default/knn/models/search.yaml
@@ -27,6 +27,10 @@ prologues:
         - {recommendation_vector: [1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], duration: 12.2}
         - {index: {_index: movies, _id: '2'}}
         - {recommendation_vector: [2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5], duration: 7.1}
+        - {index: {_index: movies, _id: '3'}}
+        - {recommendation_vector: [3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5], duration: 4.8}
+        - {index: {_index: movies, _id: '4'}}
+        - {recommendation_vector: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5], duration: 18.6}
     status: [200]
   - method: POST
     path: /_plugins/_knn/models/model-1/_train
@@ -41,6 +45,15 @@ prologues:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 4
+            encoder:
+              name: pq
+              parameters:
+                code_size: 2
     status: [200]
 epilogues:
   - path: /movies