Change VectorReaderListener to expect number array (#417)

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 · 7fcd603 · 7fcd603
1 parent 55717ed
commit 7fcd603
Show file tree

Hide file tree

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