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

* add serialization and deserialization Signed-off-by: zhichao-aws <[email protected]> * hash, equals. + UT Signed-off-by: zhichao-aws <[email protected]> * tidy Signed-off-by: zhichao-aws <[email protected]> * add test Signed-off-by: zhichao-aws <[email protected]> --------- Signed-off-by: zhichao-aws <[email protected]> (cherry picked from commit 51e6c00)
opensearch-project · Oct 11, 2023 · 3e57d7d · 3e57d7d
1 parent 0ed523b
commit 3e57d7d
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
@@ -89,6 +89,10 @@ public NeuralSparseQueryBuilder(StreamInput in) throws IOException {
         this.queryText = in.readString();
         this.modelId = in.readString();
         this.maxTokenScore = in.readOptionalFloat();
+        if (in.readBoolean()) {
+            Map<String, Float> queryTokens = in.readMap(StreamInput::readString, StreamInput::readFloat);
+            this.queryTokensSupplier = () -> queryTokens;
+        }
     }
 
     @Override
@@ -97,6 +101,12 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         out.writeString(queryText);
         out.writeString(modelId);
         out.writeOptionalFloat(maxTokenScore);
+        if (queryTokensSupplier != null && queryTokensSupplier.get() != null) {
+            out.writeBoolean(true);
+            out.writeMap(queryTokensSupplier.get(), StreamOutput::writeString, StreamOutput::writeFloat);
+        } else {
+            out.writeBoolean(false);
+        }
     }
 
     @Override
@@ -276,16 +286,25 @@ 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)
             .append(maxTokenScore, obj.maxTokenScore);
+        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).append(maxTokenScore).toHashCode();
+        HashCodeBuilder builder = new HashCodeBuilder().append(fieldName).append(queryText).append(modelId).append(maxTokenScore);
+        if (queryTokensSupplier != null) {
+            builder.append(queryTokensSupplier.get());
+        }
+        return builder.toHashCode();
     }
 
     @Override

diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java
@@ -27,6 +27,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;
@@ -294,6 +295,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() {
@@ -309,6 +327,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)
@@ -379,6 +399,24 @@ public void testHashAndEquals() {
             .boost(boost1)
             .queryName(queryName1);
 
+        // Identical to sparseEncodingQueryBuilder_baseline except non-null query tokens supplier
+        NeuralSparseQueryBuilder sparseEncodingQueryBuilder_nonNullQueryTokens = new NeuralSparseQueryBuilder().fieldName(fieldName1)
+            .queryText(queryText1)
+            .modelId(modelId1)
+            .maxTokenScore(maxTokenScore1)
+            .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)
+            .maxTokenScore(maxTokenScore1)
+            .boost(boost1)
+            .queryName(queryName1)
+            .queryTokensSupplier(() -> queryTokens2);
+
         assertEquals(sparseEncodingQueryBuilder_baseline, sparseEncodingQueryBuilder_baseline);
         assertEquals(sparseEncodingQueryBuilder_baseline.hashCode(), sparseEncodingQueryBuilder_baseline.hashCode());
 
@@ -405,6 +443,12 @@ public void testHashAndEquals() {
 
         assertNotEquals(sparseEncodingQueryBuilder_baseline, sparseEncodingQueryBuilder_diffMaxTokenScore);
         assertNotEquals(sparseEncodingQueryBuilder_baseline.hashCode(), sparseEncodingQueryBuilder_diffMaxTokenScore.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