Change VectorReaderListener to expect number array (#416)

Refactors VectorReaderListener onResponse to expect arrays of Number type from search result instead of Double type. Adds test case to confirm that it can handle Integer type. Cleans up tests in VectorReaderTest class. Signed-off-by: John Mazanec <[email protected]> (cherry picked from commit 7735351)
opensearch-project · Jun 3, 2022 · 60f50eb · 60f50eb
1 parent b83c88f
commit 60f50eb
Show file tree

Hide file tree

Showing 2 changed files with 195 additions and 97 deletions.
diff --git a/src/main/java/org/opensearch/knn/training/VectorReader.java b/src/main/java/org/opensearch/knn/training/VectorReader.java
@@ -165,9 +165,9 @@ public void onResponse(SearchResponse searchResponse) {
             List<Float[]> trainingData = new ArrayList<>();
 
             for (int i = 0; i < vectorsToAdd; i++) {
-                trainingData.add(((List<Double>) hits[i].getSourceAsMap().get(fieldName)).stream()
-                        .map(Double::floatValue)
-                        .toArray(Float[]::new));
+                trainingData.add(
+                    ((List<Number>) hits[i].getSourceAsMap().get(fieldName)).stream().map(Number::floatValue).toArray(Float[]::new)
+                );
             }
 
             this.collectedVectorCount += trainingData.size();