Repository: chishui/neural-search
Pull request #1: Add ingest processor for rewrite and modify NeuralSparseQueryBuilder
Author: login: chishui
Created: 2025-03-12T08:04:01Z | Merged: 2025-03-12T08:26:15Z

### Description
1. Add a new ingest processor to append cluster id to sparse tokens, if a sparse_encoding processor is used, this should be put after it.
2. Add logic in NeuralSparseQueryBuilder to rewrite tokens with cluster ids

create index
```
PUT {{index}}
{
  "mappings": {
    "properties": {
      "passage_text": {
        "type": "text"
      },
      "passage_embedding":{
          "type": "rank_features"
      }
    }
  }
}
```

ingest process example:
```
PUT _ingest/pipeline/ingest-pipeline
{
  "description": "This pipeline processes sparse tokens",
  "processors": [
    {
      "rewrite_token": {
        "token_field": "passage_embedding"
      },
      "remove": {
          "field": "cluster_id"
      }
    }
  ]
}
```

Bulk example:
```
POST _bulk?pipeline=ingest-pipeline
{ "index": { "_index": "{{index}}", "_id": "1" } }
{ "passage_text": "hello world", "passage_embedding": {"hello": 1.5, "world": 2.5}, "cluster_id": "cluster_1" }
{ "index": { "_index": "{{index}}", "_id": "2" } }
{ "passage_text": "apple banana", "passage_embedding": {"apple": 0.5, "banana": 1.3}, "cluster_id": "cluster_2"}
{ "index": { "_index": "{{index}}", "_id": "3" } }
{ "passage_text": "apple tree", "passage_embedding": {"apple": 3.5, "tree": 1.6}, "cluster_id": "cluster_3"}
```

Query example:
```
GET {{index}}/_search
{
    "query": {
        "neural_sparse": {
            "passage_embedding": {
                "query_tokens": {
                    "apple": 0.1,
                    "tree": 1.0
                },
                "search_cluster": true,
                "document_ratio": 0.1
            }
        }
    }
}
```

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check [here](https://github.com/opensearch-project/neural-search/blob/main/CONTRIBUTING.md#developer-certificate-of-origin).



Diff:
diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
index fad22ba95f..0b84de9190 100644
--- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
+++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
@@ -14,6 +14,7 @@
 import java.util.Optional;
 import java.util.function.Supplier;
 
+import org.opensearch.neuralsearch.processor.RewriteTokenProcessor;
 import org.opensearch.transport.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.service.ClusterService;
@@ -133,7 +134,9 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
             TextImageEmbeddingProcessor.TYPE,
             new TextImageEmbeddingProcessorFactory(clientAccessor, parameters.env, parameters.ingestService.getClusterService()),
             TextChunkingProcessor.TYPE,
-            new TextChunkingProcessorFactory(parameters.env, parameters.ingestService.getClusterService(), parameters.analysisRegistry)
+            new TextChunkingProcessorFactory(parameters.env, parameters.ingestService.getClusterService(), parameters.analysisRegistry),
+            RewriteTokenProcessor.TYPE,
+            new RewriteTokenProcessor.Factory()
         );
     }
 
diff --git a/src/main/java/org/opensearch/neuralsearch/processor/RewriteTokenProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/RewriteTokenProcessor.java
new file mode 100644
index 0000000000..146b325d28
--- /dev/null
+++ b/src/main/java/org/opensearch/neuralsearch/processor/RewriteTokenProcessor.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.processor;
+
+import org.opensearch.ingest.AbstractProcessor;
+import org.opensearch.ingest.IngestDocument;
+import org.opensearch.ingest.Processor;
+import org.opensearch.neuralsearch.processor.util.ClusterUtils;
+
+import java.util.Map;
+
+import static org.opensearch.ingest.ConfigurationUtils.readStringProperty;
+
+/**
+ * Processor for rewriting token field to new token field
+ */
+public class RewriteTokenProcessor extends AbstractProcessor {
+    public static final String TYPE = "rewrite_token";
+    public static final String TOKEN_FIELD_KEY = "token_field";
+    public static final String CLUSTER_ID = "cluster_id";
+    private String tokenField;
+
+    protected RewriteTokenProcessor(String tag, String description, String tokenField) {
+        super(tag, description);
+        this.tokenField = tokenField;
+    }
+
+    @Override
+    public IngestDocument execute(IngestDocument ingestDocument) throws Exception {
+        Map<String, Float> tokens = ingestDocument.getFieldValue(tokenField, Map.class);
+        String clusterId = ingestDocument.getFieldValue(CLUSTER_ID, String.class);
+        Map<String, Float> newTokens = tokens.entrySet()
+            .stream()
+            .collect(java.util.stream.Collectors.toMap(e -> ClusterUtils.constructNewToken(e.getKey(), clusterId), Map.Entry::getValue));
+        ingestDocument.setFieldValue(tokenField, newTokens);
+        return ingestDocument;
+    }
+
+    @Override
+    public String getType() {
+        return TYPE;
+    }
+
+    public static final class Factory implements Processor.Factory {
+
+        @Override
+        public Processor create(
+            Map<String, Processor.Factory> processorFactories,
+            String tag,
+            String description,
+            Map<String, Object> config
+        ) throws Exception {
+            String tokenField = readStringProperty(TYPE, tag, config, TOKEN_FIELD_KEY);
+            return new RewriteTokenProcessor(tag, description, tokenField);
+        }
+    }
+}
diff --git a/src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java b/src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java
new file mode 100644
index 0000000000..226fa8dba9
--- /dev/null
+++ b/src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.processor.util;
+
+/**
+ * Utility class for cluster related operations
+ */
+public class ClusterUtils {
+    public static String constructNewToken(String token, String clusterId) {
+        return token + "_" + clusterId;
+    }
+}
diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
index 2e27fae58c..7109744f3c 100644
--- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
+++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java
@@ -5,6 +5,7 @@
 package org.opensearch.neuralsearch.query;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -16,11 +17,13 @@
 import org.apache.commons.lang.StringUtils;
 import org.apache.commons.lang.builder.EqualsBuilder;
 import org.apache.commons.lang.builder.HashCodeBuilder;
+import org.apache.commons.lang3.BooleanUtils;
 import org.apache.lucene.document.FeatureField;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.Version;
+import org.opensearch.neuralsearch.processor.util.ClusterUtils;
 import org.opensearch.transport.client.Client;
 import org.opensearch.common.SetOnce;
 import org.opensearch.common.collect.Tuple;
@@ -70,6 +73,8 @@ public class NeuralSparseQueryBuilder extends AbstractQueryBuilder<NeuralSparseQ
     static final ParseField QUERY_TOKENS_FIELD = new ParseField("query_tokens");
     @VisibleForTesting
     static final ParseField MODEL_ID_FIELD = new ParseField("model_id");
+    static final ParseField SEARCH_CLUSTER_FIELD = new ParseField("search_cluster");
+    static final ParseField DOCUMENT_RATIO_FIELD = new ParseField("document_ratio");
     // 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
@@ -80,6 +85,8 @@ public class NeuralSparseQueryBuilder extends AbstractQueryBuilder<NeuralSparseQ
     private String queryText;
     private String modelId;
     private Float maxTokenScore;
+    private Boolean searchCluster;
+    private Float documentRatio;
     private Supplier<Map<String, Float>> queryTokensSupplier;
     // A field that for neural_sparse_two_phase_processor, if twoPhaseSharedQueryToken is not null,
     // it means it's origin NeuralSparseQueryBuilder and should split the low score tokens form itself then put it into
@@ -127,6 +134,8 @@ public NeuralSparseQueryBuilder(StreamInput in) throws IOException {
         if (StringUtils.EMPTY.equals(this.modelId)) {
             this.modelId = null;
         }
+        this.searchCluster = in.readOptionalBoolean();
+        this.documentRatio = in.readOptionalFloat();
     }
 
     /**
@@ -142,7 +151,9 @@ public NeuralSparseQueryBuilder getCopyNeuralSparseQueryBuilderForTwoPhase(float
             .queryText(this.queryText)
             .modelId(this.modelId)
             .maxTokenScore(this.maxTokenScore)
-            .twoPhasePruneRatio(-1f * pruneRatio);
+            .twoPhasePruneRatio(-1f * pruneRatio)
+            .searchCluster(this.searchCluster)
+            .documentRatio(this.documentRatio);
         if (Objects.nonNull(this.queryTokensSupplier)) {
             Map<String, Float> tokens = queryTokensSupplier.get();
             // Splitting tokens based on a threshold value: tokens greater than the threshold are stored in v1,
@@ -175,6 +186,8 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
+        out.writeOptionalBoolean(this.searchCluster);
+        out.writeOptionalFloat(this.documentRatio);
     }
 
     @Override
@@ -193,6 +206,10 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws
         if (Objects.nonNull(queryTokensSupplier) && Objects.nonNull(queryTokensSupplier.get())) {
             xContentBuilder.field(QUERY_TOKENS_FIELD.getPreferredName(), queryTokensSupplier.get());
         }
+        xContentBuilder.field(SEARCH_CLUSTER_FIELD.getPreferredName(), searchCluster);
+        if (Objects.nonNull(documentRatio)) {
+            xContentBuilder.field(DOCUMENT_RATIO_FIELD.getPreferredName(), documentRatio);
+        }
         printBoostAndQueryName(xContentBuilder);
         xContentBuilder.endObject();
         xContentBuilder.endObject();
@@ -297,6 +314,10 @@ private static void parseQueryParams(XContentParser parser, NeuralSparseQueryBui
                     sparseEncodingQueryBuilder.modelId(parser.text());
                 } else if (MAX_TOKEN_SCORE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     sparseEncodingQueryBuilder.maxTokenScore(parser.floatValue());
+                } else if (SEARCH_CLUSTER_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    sparseEncodingQueryBuilder.searchCluster(parser.booleanValue());
+                } else if (DOCUMENT_RATIO_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    sparseEncodingQueryBuilder.documentRatio(parser.floatValue());
                 } else {
                     throw new ParsingException(
                         parser.getTokenLocation(),
@@ -315,6 +336,27 @@ private static void parseQueryParams(XContentParser parser, NeuralSparseQueryBui
         }
     }
 
+    private List<String> getClusterIds(Map<String, Float> queryTokens) {
+        // step 1: transform tokens to dense
+        // step 2: call cluster service to get top clusters with ratio
+        return Arrays.asList("cluster_1", "cluster_2", "cluster_3");
+    }
+
+    private Map<String, Float> generateNewQueryTokensBasedOnClusters(Map<String, Float> queryTokens) {
+        if (!BooleanUtils.isTrue(this.searchCluster)) {
+            return queryTokens;
+        }
+        final List<String> clusterIds = getClusterIds(queryTokens);
+        Map<String, Float> newQueryTokens = new HashMap<>();
+        for (Map.Entry<String, Float> entry : queryTokens.entrySet()) {
+            for (String clusterId : clusterIds) {
+                String newToken = ClusterUtils.constructNewToken(entry.getKey(), clusterId);
+                newQueryTokens.put(newToken, entry.getValue());
+            }
+        }
+        return newQueryTokens;
+    }
+
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
         // We need to inference the sentence to get the queryTokens. The logic is similar to NeuralQueryBuilder
@@ -367,7 +409,7 @@ private BiConsumer<Client, ActionListener<?>> getModelInferenceAsync(SetOnce<Map
     protected Query doToQuery(QueryShardContext context) throws IOException {
         final MappedFieldType ft = context.fieldMapper(fieldName);
         validateFieldType(ft);
-        Map<String, Float> queryTokens = queryTokensSupplier.get();
+        Map<String, Float> queryTokens = generateNewQueryTokensBasedOnClusters(queryTokensSupplier.get());
         if (Objects.isNull(queryTokens)) {
             throw new IllegalArgumentException("Query tokens cannot be null.");
         }
@@ -417,7 +459,9 @@ protected boolean doEquals(NeuralSparseQueryBuilder obj) {
             .append(modelId, obj.modelId)
             .append(maxTokenScore, obj.maxTokenScore)
             .append(twoPhasePruneRatio, obj.twoPhasePruneRatio)
-            .append(twoPhaseSharedQueryToken, obj.twoPhaseSharedQueryToken);
+            .append(twoPhaseSharedQueryToken, obj.twoPhaseSharedQueryToken)
+            .append(searchCluster, obj.searchCluster)
+            .append(documentRatio, obj.documentRatio);
         if (Objects.nonNull(queryTokensSupplier)) {
             equalsBuilder.append(queryTokensSupplier.get(), obj.queryTokensSupplier.get());
         }
@@ -431,7 +475,9 @@ protected int doHashCode() {
             .append(modelId)
             .append(maxTokenScore)
             .append(twoPhasePruneRatio)
-            .append(twoPhaseSharedQueryToken);
+            .append(twoPhaseSharedQueryToken)
+            .append(searchCluster)
+            .append(documentRatio);
         if (Objects.nonNull(queryTokensSupplier)) {
             builder.append(queryTokensSupplier.get());
         }


Reviews:
user: login: yuye-aws
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T08:10:44Z
user: login: chishui
state: COMMENTED
author_association: OWNER
submitted_at: 2025-03-12T08:16:13Z
user: login: yuye-aws
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T08:20:11Z

Review comments:
diff_hunk: @@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.processor.util;
+
+/**
+ * Utility class for cluster related operations
+ */
+public class ClusterUtils {
path: src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java
user: login: yuye-aws
body: Just wondering if cluster can be mixed with the OpenSearch keywords. My code is ANN utils. Or maybe KMeansClusterUtils?
created_at: 2025-03-12T08:08:16Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.processor.util;
+
+/**
+ * Utility class for cluster related operations
+ */
+public class ClusterUtils {
path: src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java
user: login: chishui
body: What about DocumentClusterUtils to not tied to algorithm name?
created_at: 2025-03-12T08:16:13Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.processor.util;
+
+/**
+ * Utility class for cluster related operations
+ */
+public class ClusterUtils {
path: src/main/java/org/opensearch/neuralsearch/processor/util/ClusterUtils.java
user: login: yuye-aws
body: Agreed
created_at: 2025-03-12T08:20:11Z
author_association: COLLABORATOR

PR comments:
user: login: yuye-aws
created_at: 2025-03-12T08:20:46Z
body: When searching, I meet with the following error. Can you please share the index mapping so that I can quickly test:

```
{
  "error": {
    "root_cause": [
      {
        "type": "query_shard_exception",
        "reason": "failed to create query: [neural_sparse] query only works on [rank_features] fields",
        "index": "{{index}}",
        "index_uuid": "HCNgB0LFSVmNKt5cFIRUJA"
      }
    ],
    "type": "search_phase_execution_exception",
    "reason": "all shards failed",
    "phase": "query",
    "grouped": true,
    "failed_shards": [
      {
        "shard": 0,
        "index": "{{index}}",
        "node": "ufb7A7DQTXGZF5CbhZegvA",
        "reason": {
          "type": "query_shard_exception",
          "reason": "failed to create query: [neural_sparse] query only works on [rank_features] fields",
          "index": "{{index}}",
          "index_uuid": "HCNgB0LFSVmNKt5cFIRUJA",
          "caused_by": {
            "type": "illegal_argument_exception",
            "reason": "[neural_sparse] query only works on [rank_features] fields"
          }
        }
      }
    ]
  },
  "status": 400
}
```
author_association: COLLABORATOR
user: login: chishui
created_at: 2025-03-12T09:33:30Z
body: test
author_association: OWNER