-
Notifications
You must be signed in to change notification settings - Fork 83
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Encapsulate KNNQueryBuilder creation within NeuralKNNQueryBuilder
Signed-off-by: Junqiu Lei <junqiu@amazon.com>
- Loading branch information
1 parent
7dc84b5
commit ebfc232
Showing
8 changed files
with
430 additions
and
32 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
68 changes: 68 additions & 0 deletions
68
src/main/java/org/opensearch/neuralsearch/query/NeuralKNNQuery.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
| package org.opensearch.neuralsearch.query; | ||
|
|
||
| import lombok.Getter; | ||
| import org.apache.lucene.index.IndexReader; | ||
| import org.apache.lucene.search.IndexSearcher; | ||
| import org.apache.lucene.search.Query; | ||
| import org.apache.lucene.search.QueryVisitor; | ||
| import org.apache.lucene.search.ScoreMode; | ||
| import org.apache.lucene.search.Weight; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Objects; | ||
|
|
||
| /** | ||
| * Wraps KNN Lucene query to support neural search extensions. | ||
| * Delegates core operations to the underlying KNN query. | ||
| */ | ||
| @Getter | ||
| public class NeuralKNNQuery extends Query { | ||
| private final Query knnQuery; | ||
|
|
||
| public NeuralKNNQuery(Query knnQuery) { | ||
| this.knnQuery = knnQuery; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString(String field) { | ||
| return knnQuery.toString(field); | ||
| } | ||
|
|
||
| @Override | ||
| public void visit(QueryVisitor visitor) { | ||
| // Delegate the visitor to the underlying KNN query | ||
| knnQuery.visit(visitor); | ||
| } | ||
|
|
||
| @Override | ||
| public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { | ||
| // Delegate weight creation to the underlying KNN query | ||
| return knnQuery.createWeight(searcher, scoreMode, boost); | ||
| } | ||
|
|
||
| @Override | ||
| public Query rewrite(IndexReader reader) throws IOException { | ||
| Query rewritten = knnQuery.rewrite(reader); | ||
| if (rewritten == knnQuery) { | ||
| return this; | ||
| } | ||
| return new NeuralKNNQuery(rewritten); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object other) { | ||
| if (this == other) return true; | ||
| if (other == null || getClass() != other.getClass()) return false; | ||
| NeuralKNNQuery that = (NeuralKNNQuery) other; | ||
| return Objects.equals(knnQuery, that.knnQuery); | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return Objects.hash(knnQuery); | ||
| } | ||
| } |
167 changes: 167 additions & 0 deletions
167
src/main/java/org/opensearch/neuralsearch/query/NeuralKNNQueryBuilder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,167 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
| package org.opensearch.neuralsearch.query; | ||
|
|
||
| import lombok.Getter; | ||
| import org.apache.lucene.search.Query; | ||
| import org.opensearch.core.common.io.stream.StreamOutput; | ||
| import org.opensearch.core.xcontent.XContentBuilder; | ||
| import org.opensearch.index.query.AbstractQueryBuilder; | ||
| import org.opensearch.index.query.QueryBuilder; | ||
| import org.opensearch.index.query.QueryRewriteContext; | ||
| import org.opensearch.index.query.QueryShardContext; | ||
| import org.opensearch.knn.index.query.KNNQueryBuilder; | ||
| import org.opensearch.knn.index.query.rescore.RescoreContext; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Map; | ||
| import java.util.Objects; | ||
|
|
||
| /** | ||
| * NeuralKNNQueryBuilder wraps KNNQueryBuilder to: | ||
| * 1. Isolate KNN plugin API changes to a single location | ||
| * 2. Allow extension with neural-search-specific information (e.g., query text) | ||
| */ | ||
|
|
||
| @Getter | ||
| public class NeuralKNNQueryBuilder extends AbstractQueryBuilder<NeuralKNNQueryBuilder> { | ||
| private final KNNQueryBuilder knnQueryBuilder; | ||
|
|
||
| /** | ||
| * Creates a new builder instance. | ||
| */ | ||
| public static Builder builder() { | ||
| return new Builder(); | ||
| } | ||
|
|
||
| public String fieldName() { | ||
| return knnQueryBuilder.fieldName(); | ||
| } | ||
|
|
||
| public int k() { | ||
| return knnQueryBuilder.getK(); | ||
| } | ||
|
|
||
| /** | ||
| * Builder for NeuralKNNQueryBuilder. | ||
| */ | ||
| public static class Builder { | ||
| private String fieldName; | ||
| private float[] vector; | ||
| private Integer k; | ||
| private QueryBuilder filter; | ||
| private Float maxDistance; | ||
| private Float minScore; | ||
| private Boolean expandNested; | ||
| private Map<String, ?> methodParameters; | ||
| private RescoreContext rescoreContext; | ||
|
|
||
| private Builder() {} | ||
|
|
||
| public Builder fieldName(String fieldName) { | ||
| this.fieldName = fieldName; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder vector(float[] vector) { | ||
| this.vector = vector; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder k(Integer k) { | ||
| this.k = k; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder filter(QueryBuilder filter) { | ||
| this.filter = filter; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder maxDistance(Float maxDistance) { | ||
| this.maxDistance = maxDistance; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder minScore(Float minScore) { | ||
| this.minScore = minScore; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder expandNested(Boolean expandNested) { | ||
| this.expandNested = expandNested; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder methodParameters(Map<String, ?> methodParameters) { | ||
| this.methodParameters = methodParameters; | ||
| return this; | ||
| } | ||
|
|
||
| public Builder rescoreContext(RescoreContext rescoreContext) { | ||
| this.rescoreContext = rescoreContext; | ||
| return this; | ||
| } | ||
|
|
||
| public NeuralKNNQueryBuilder build() { | ||
| KNNQueryBuilder knnBuilder = KNNQueryBuilder.builder() | ||
| .fieldName(fieldName) | ||
| .vector(vector) | ||
| .k(k) | ||
| .filter(filter) | ||
| .maxDistance(maxDistance) | ||
| .minScore(minScore) | ||
| .expandNested(expandNested) | ||
| .methodParameters(methodParameters) | ||
| .rescoreContext(rescoreContext) | ||
| .build(); | ||
| return new NeuralKNNQueryBuilder(knnBuilder); | ||
| } | ||
| } | ||
|
|
||
| private NeuralKNNQueryBuilder(KNNQueryBuilder knnQueryBuilder) { | ||
| this.knnQueryBuilder = knnQueryBuilder; | ||
| } | ||
|
|
||
| @Override | ||
| public void doWriteTo(StreamOutput out) throws IOException { | ||
| knnQueryBuilder.writeTo(out); | ||
| } | ||
|
|
||
| @Override | ||
| protected void doXContent(XContentBuilder builder, Params params) throws IOException { | ||
| knnQueryBuilder.toXContent(builder, params); | ||
| } | ||
|
|
||
| @Override | ||
| protected QueryBuilder doRewrite(QueryRewriteContext context) throws IOException { | ||
| QueryBuilder rewritten = knnQueryBuilder.rewrite(context); | ||
| if (rewritten == knnQueryBuilder) { | ||
| return this; | ||
| } | ||
| return new NeuralKNNQueryBuilder((KNNQueryBuilder) rewritten); | ||
| } | ||
|
|
||
| @Override | ||
| protected Query doToQuery(QueryShardContext context) throws IOException { | ||
| Query knnQuery = knnQueryBuilder.toQuery(context); | ||
| return new NeuralKNNQuery(knnQuery); | ||
| } | ||
|
|
||
| @Override | ||
| protected boolean doEquals(NeuralKNNQueryBuilder other) { | ||
| return Objects.equals(knnQueryBuilder, other.knnQueryBuilder); | ||
| } | ||
|
|
||
| @Override | ||
| protected int doHashCode() { | ||
| return Objects.hash(knnQueryBuilder); | ||
| } | ||
|
|
||
| @Override | ||
| public String getWriteableName() { | ||
| return knnQueryBuilder.getWriteableName(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.