[Backport main manually][bug fix] Fix async actions are left in neural_sparse query (#438)

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 ### Enhancements
 ### Bug Fixes
+Fix async actions are left in neural_sparse query ([438](https://github.com/opensearch-project/neural-search/pull/438))
 ### Infrastructure
 ### Documentation
 ### Maintenance

src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java

@@ -84,13 +84,23 @@ public NeuralSparseQueryBuilder(StreamInput in) throws IOException {
         this.fieldName = in.readString();
         this.queryText = in.readString();
         this.modelId = in.readString();
+        if (in.readBoolean()) {
+            Map<String, Float> queryTokens = in.readMap(StreamInput::readString, StreamInput::readFloat);
+            this.queryTokensSupplier = () -> queryTokens;
+        }
     }
 
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         out.writeString(fieldName);
         out.writeString(queryText);
         out.writeString(modelId);
+        if (queryTokensSupplier != null && queryTokensSupplier.get() != null) {
+            out.writeBoolean(true);
+            out.writeMap(queryTokensSupplier.get(), StreamOutput::writeString, StreamOutput::writeFloat);
+        } else {
+            out.writeBoolean(false);
+        }
     }
 
     @Override
@@ -257,15 +267,24 @@ private static void validateQueryTokens(Map<String, Float> queryTokens) {
     protected boolean doEquals(NeuralSparseQueryBuilder obj) {
         if (this == obj) return true;
         if (obj == null || getClass() != obj.getClass()) return false;
+        if (queryTokensSupplier == null && obj.queryTokensSupplier != null) return false;
+        if (queryTokensSupplier != null && obj.queryTokensSupplier == null) return false;
         EqualsBuilder equalsBuilder = new EqualsBuilder().append(fieldName, obj.fieldName)
             .append(queryText, obj.queryText)
             .append(modelId, obj.modelId);
+        if (queryTokensSupplier != null) {
+            equalsBuilder.append(queryTokensSupplier.get(), obj.queryTokensSupplier.get());
+        }
         return equalsBuilder.isEquals();
     }
 
     @Override
     protected int doHashCode() {
-        return new HashCodeBuilder().append(fieldName).append(queryText).append(modelId).toHashCode();
+        HashCodeBuilder builder = new HashCodeBuilder().append(fieldName).append(queryText).append(modelId);
+        if (queryTokensSupplier != null) {
+            builder.append(queryTokensSupplier.get());
+        }
+        return builder.toHashCode();
     }
 
     @Override

src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java

@@ -26,6 +26,7 @@
 import lombok.SneakyThrows;
 
 import org.opensearch.client.Client;
+import org.opensearch.common.SetOnce;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.action.ActionListener;
@@ -262,6 +263,23 @@ public void testStreams() {
 
         NeuralSparseQueryBuilder copy = new NeuralSparseQueryBuilder(filterStreamInput);
         assertEquals(original, copy);
+
+        SetOnce<Map<String, Float>> queryTokensSetOnce = new SetOnce<>();
+        queryTokensSetOnce.set(Map.of("hello", 1.0f, "world", 2.0f));
+        original.queryTokensSupplier(queryTokensSetOnce::get);
+
+        streamOutput = new BytesStreamOutput();
+        original.writeTo(streamOutput);
+
+        filterStreamInput = new NamedWriteableAwareStreamInput(
+            streamOutput.bytes().streamInput(),
+            new NamedWriteableRegistry(
+                List.of(new NamedWriteableRegistry.Entry(QueryBuilder.class, MatchAllQueryBuilder.NAME, MatchAllQueryBuilder::new))
+            )
+        );
+
+        copy = new NeuralSparseQueryBuilder(filterStreamInput);
+        assertEquals(original, copy);
     }
 
     public void testHashAndEquals() {
@@ -275,6 +293,8 @@ public void testHashAndEquals() {
         float boost2 = 3.8f;
         String queryName1 = "query-1";
         String queryName2 = "query-2";
+        Map<String, Float> queryTokens1 = Map.of("hello", 1.0f, "world", 2.0f);
+        Map<String, Float> queryTokens2 = Map.of("hello", 1.0f, "world", 2.2f);
 
         NeuralSparseQueryBuilder sparseEncodingQueryBuilder_baseline = new NeuralSparseQueryBuilder().fieldName(fieldName1)
             .queryText(queryText1)
@@ -329,6 +349,22 @@ public void testHashAndEquals() {
             .boost(boost1)
             .queryName(queryName2);
 
+        // Identical to sparseEncodingQueryBuilder_baseline except non-null query tokens supplier
+        NeuralSparseQueryBuilder sparseEncodingQueryBuilder_nonNullQueryTokens = new NeuralSparseQueryBuilder().fieldName(fieldName1)
+            .queryText(queryText1)
+            .modelId(modelId1)
+            .boost(boost1)
+            .queryName(queryName1)
+            .queryTokensSupplier(() -> queryTokens1);
+
+        // Identical to sparseEncodingQueryBuilder_baseline except non-null query tokens supplier
+        NeuralSparseQueryBuilder sparseEncodingQueryBuilder_diffQueryTokens = new NeuralSparseQueryBuilder().fieldName(fieldName1)
+            .queryText(queryText1)
+            .modelId(modelId1)
+            .boost(boost1)
+            .queryName(queryName1)
+            .queryTokensSupplier(() -> queryTokens2);
+
         assertEquals(sparseEncodingQueryBuilder_baseline, sparseEncodingQueryBuilder_baseline);
         assertEquals(sparseEncodingQueryBuilder_baseline.hashCode(), sparseEncodingQueryBuilder_baseline.hashCode());
 
@@ -352,6 +388,12 @@ public void testHashAndEquals() {
 
         assertNotEquals(sparseEncodingQueryBuilder_baseline, sparseEncodingQueryBuilder_diffQueryName);
         assertNotEquals(sparseEncodingQueryBuilder_baseline.hashCode(), sparseEncodingQueryBuilder_diffQueryName.hashCode());
+
+        assertNotEquals(sparseEncodingQueryBuilder_baseline, sparseEncodingQueryBuilder_nonNullQueryTokens);
+        assertNotEquals(sparseEncodingQueryBuilder_baseline.hashCode(), sparseEncodingQueryBuilder_nonNullQueryTokens.hashCode());
+
+        assertNotEquals(sparseEncodingQueryBuilder_nonNullQueryTokens, sparseEncodingQueryBuilder_diffQueryTokens);
+        assertNotEquals(sparseEncodingQueryBuilder_nonNullQueryTokens.hashCode(), sparseEncodingQueryBuilder_diffQueryTokens.hashCode());
     }
 
     @SneakyThrows