From af0e58ca56f1c568816b86810baf1bf957b30459 Mon Sep 17 00:00:00 2001 From: Vijayan Balasubramanian Date: Sat, 4 Jan 2025 23:34:35 -0800 Subject: [PATCH] in middle Signed-off-by: Vijayan Balasubramanian --- .../opensearch/knn/bwc/ScriptScoringIT.java | 2 +- .../opensearch/knn/bwc/ScriptScoringIT.java | 61 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java index c9fbd82d0b..ec8378321f 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java @@ -56,7 +56,7 @@ public void testKNNCosineScriptScore() throws Exception { float[] indexVector2 = { 8.1f, 9.1f, 0.3f }; float[] queryVector = { 3.0f, 4.0f, 5.5f }; if (isRunningAgainstOldCluster()) { - createKnnIndex(testIndex, createKNNDefaultScriptScoreSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS)); + createKnnIndex(testIndex, createKNNDefaultScriptScoreSettings(), createKnnIndexMapping(TEST_FIELD, 3)); addKnnDoc(testIndex, "1", TEST_FIELD, indexVector1); validateScore(1, queryVector, new float[] { VectorUtil.cosine(queryVector, indexVector1) }); } else { diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java index 1902571120..8e343df3de 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ScriptScoringIT.java @@ -5,8 +5,13 @@ package org.opensearch.knn.bwc; +import org.apache.hc.core5.http.io.entity.EntityUtils; +import org.apache.lucene.util.VectorUtil; +import org.opensearch.client.Response; import org.opensearch.knn.index.SpaceType; +import java.util.List; + import static org.opensearch.knn.TestUtils.NODES_BWC_CLUSTER; public class ScriptScoringIT extends AbstractRollingUpgradeTestCase { @@ -54,4 +59,60 @@ public void validateKNNL2ScriptScoreOnUpgrade(int totalDocsCount, int docId) thr validateKNNScriptScoreSearch(testIndex, TEST_FIELD, DIMENSIONS, totalDocsCount, K, SpaceType.L2); } + // KNN script scoring for space_type "cosine" + public void testKNNCosineScriptScore() throws Exception { + float[] indexVector1 = { 1.1f, 2.1f, 3.3f }; + float[] indexVector2 = { 8.1f, 9.1f, 10.3f }; + float[] indexVector3 = { 9.1f, 10.1f, 11.3f }; + float[] indexVector4 = { 10.1f, 11.1f, 12.3f }; + float[] queryVector = { 3.0f, 4.0f, 13.5f }; + waitForClusterHealthGreen(NODES_BWC_CLUSTER); + int k = 10; + switch (getClusterType()) { + case OLD: + createKnnIndex(testIndex, createKNNDefaultScriptScoreSettings(), createKnnIndexMapping(TEST_FIELD, 3)); + addKnnDoc(testIndex, "1", TEST_FIELD, indexVector1); + validateScore(k, queryVector, List.of(VectorUtil.cosine(queryVector, indexVector1))); + break; + case MIXED: + if (isFirstMixedRound()) { + addKnnDoc(testIndex, "2", TEST_FIELD, indexVector2); + validateScore( + k, + queryVector, + List.of(VectorUtil.cosine(queryVector, indexVector1), VectorUtil.cosine(queryVector, indexVector2)) + ); + } else { + addKnnDoc(testIndex, "3", TEST_FIELD, indexVector3); + validateScore( + k, + queryVector, + List.of( + VectorUtil.cosine(queryVector, indexVector1), + VectorUtil.cosine(queryVector, indexVector2), + VectorUtil.cosine(queryVector, indexVector3) + ) + ); + } + break; + case UPGRADED: + addKnnDoc(testIndex, "3", TEST_FIELD, indexVector3); + validateScore( + k, + queryVector, + List.of( + VectorUtil.cosine(queryVector, indexVector1), + VectorUtil.cosine(queryVector, indexVector2), + VectorUtil.cosine(queryVector, indexVector3), + VectorUtil.cosine(queryVector, indexVector4) + ) + ); + } + } + + private void validateScore(int k, float[] queryVector, List expectedScores) throws Exception { + final Response responseBody = executeKNNScriptScoreRequest(testIndex, TEST_FIELD, k, SpaceType.COSINESIMIL, queryVector); + List actualScores = parseSearchResponseScore(EntityUtils.toString(responseBody.getEntity()), TEST_FIELD); + assertEquals(expectedScores, actualScores); + } }