From 6a0bb56132e264ff91a0503308836d191b6b563e Mon Sep 17 00:00:00 2001
From: zhichao-aws <zhichaog@amazon.com>
Date: Mon, 30 Oct 2023 10:39:51 +0800
Subject: [PATCH] deprecate max_token_score

Signed-off-by: zhichao-aws <zhichaog@amazon.com>
---
 .../query/NeuralSparseQueryBuilder.java       |  7 +++--
 .../query/NeuralSparseQueryBuilderTests.java  | 26 +++++++++++++++++++
 .../query/NeuralSparseQueryIT.java            | 11 +-------
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
index 2373ad060..1274107b7 100644
--- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
+++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
@@ -61,8 +61,10 @@ public class NeuralSparseQueryBuilder extends AbstractQueryBuilder<NeuralSparseQ
     static final ParseField QUERY_TEXT_FIELD = new ParseField("query_text");
     @VisibleForTesting
     static final ParseField MODEL_ID_FIELD = new ParseField("model_id");
+    // We use max_token_score field to help WAND scorer prune query clause in lucene 9.7. But in lucene 9.8 the inner
+    // logics change, this field is not needed any more.
     @VisibleForTesting
-    static final ParseField MAX_TOKEN_SCORE_FIELD = new ParseField("max_token_score");
+    static final ParseField MAX_TOKEN_SCORE_FIELD = new ParseField("max_token_score").withAllDeprecated();
 
     private static MLCommonsClientAccessor ML_CLIENT;
 
@@ -163,9 +165,6 @@ public static NeuralSparseQueryBuilder fromXContent(XContentParser parser) throw
             sparseEncodingQueryBuilder.modelId(),
             String.format(Locale.ROOT, "%s field must be provided for [%s] query", MODEL_ID_FIELD.getPreferredName(), NAME)
         );
-        if (sparseEncodingQueryBuilder.maxTokenScore != null && sparseEncodingQueryBuilder.maxTokenScore <= 0) {
-            throw new IllegalArgumentException(MAX_TOKEN_SCORE_FIELD.getPreferredName() + " must be larger than 0.");
-        }
 
         return sparseEncodingQueryBuilder;
     }
diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java
index a50ab4fb8..6588b8c77 100644
--- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java
+++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java
@@ -117,6 +117,32 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() {
         assertEquals(QUERY_NAME, sparseEncodingQueryBuilder.queryName());
     }
 
+    @SneakyThrows
+    public void testFromXContent_whenBuiltWithMaxTokenScore_thenThrowWarning() {
+        /*
+          {
+              "VECTOR_FIELD": {
+                "query_text": "string",
+                "model_id": "string",
+                "max_token_score": 123.0
+              }
+          }
+        */
+        XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()
+                .startObject()
+                .startObject(FIELD_NAME)
+                .field(QUERY_TEXT_FIELD.getPreferredName(), QUERY_TEXT)
+                .field(MODEL_ID_FIELD.getPreferredName(), MODEL_ID)
+                .field(MAX_TOKEN_SCORE_FIELD.getPreferredName(), MAX_TOKEN_SCORE)
+                .endObject()
+                .endObject();
+
+        XContentParser contentParser = createParser(xContentBuilder);
+        contentParser.nextToken();
+        NeuralSparseQueryBuilder sparseEncodingQueryBuilder = NeuralSparseQueryBuilder.fromXContent(contentParser);
+        assertWarnings("Deprecated field [max_token_score] used, this field is unused and will be removed entirely");
+    }
+
     @SneakyThrows
     public void testFromXContent_whenBuildWithMultipleRootFields_thenFail() {
         /*
diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java
index 672ab2940..12bd1c1cb 100644
--- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java
+++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java
@@ -106,16 +106,7 @@ public void testBasicQueryWithMaxTokenScore() {
         Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
 
         assertEquals("1", firstInnerHit.get("_id"));
-        Map<String, Float> queryTokens = runSparseModelInference(modelId, TEST_QUERY_TEXT);
-        float expectedScore = 0f;
-        for (Map.Entry<String, Float> entry : queryTokens.entrySet()) {
-            if (testRankFeaturesDoc.containsKey(entry.getKey())) {
-                expectedScore += entry.getValue() * Math.min(
-                    getFeatureFieldCompressedNumber(testRankFeaturesDoc.get(entry.getKey())),
-                    maxTokenScore
-                );
-            }
-        }
+        float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT);
         assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA);
     }