Support Radial Search (#1166)

* Generate the ml namespace (#1158) * Generate ml.register_model_group Signed-off-by: Thomas Farr <[email protected]> * Start neural search sample Signed-off-by: Thomas Farr <[email protected]> * Re-generate ShardStatistics Signed-off-by: Thomas Farr <[email protected]> * Re-generate ShardFailure Signed-off-by: Thomas Farr <[email protected]> * Re-generate Result Signed-off-by: Thomas Farr <[email protected]> * Re-generate WriteResponseBase Signed-off-by: Thomas Farr <[email protected]> * Generate ml.delete_model_group Signed-off-by: Thomas Farr <[email protected]> * Generate ml.register_model Signed-off-by: Thomas Farr <[email protected]> * Exclude legacy license from ml namespace Signed-off-by: Thomas Farr <[email protected]> * Generate ml.get_task Signed-off-by: Thomas Farr <[email protected]> * Generate ml.delete_task Signed-off-by: Thomas Farr <[email protected]> * Generate ml.delete_model Signed-off-by: Thomas Farr <[email protected]> * Generate ml.deploy_model Signed-off-by: Thomas Farr <[email protected]> * Generate ml.undeploy_model Signed-off-by: Thomas Farr <[email protected]> * Complete neural search sample Signed-off-by: Thomas Farr <[email protected]> * Generate ml.get_model Signed-off-by: Thomas Farr <[email protected]> * Add changelog entry Signed-off-by: Thomas Farr <[email protected]> * note Signed-off-by: Thomas Farr <[email protected]> * Fix tests Signed-off-by: Thomas Farr <[email protected]> --------- Signed-off-by: Thomas Farr <[email protected]> Signed-off-by: Alex Keeler <[email protected]> * Fix copy-paste mistake in NeuralSearch sample (#1161) Signed-off-by: Thomas Farr <[email protected]> Signed-off-by: Alex Keeler <[email protected]> * Support Radial Search Add minScore, maxDistance parameters to KnnQuery in order to support Radial Search, which was introduced in OpenSearch 2.14 https://opensearch.org/docs/latest/search-plugins/knn/radial-search-knn/ Signed-off-by: Alex Keeler <[email protected]> * Update CHANGELOG.md Signed-off-by: alex-keeler <[email protected]> Signed-off-by: Alex Keeler <[email protected]> * Update changelog post releasing 2.14.0 (#1162) (#1167) Signed-off-by: Thomas Farr <[email protected]> (cherry picked from commit 2a362a6) * Reduce required release approvals (#1168) Signed-off-by: Thomas Farr <[email protected]> Signed-off-by: Alex Keeler <[email protected]> --------- Signed-off-by: Thomas Farr <[email protected]> Signed-off-by: Alex Keeler <[email protected]> Signed-off-by: alex-keeler <[email protected]> Co-authored-by: Thomas Farr <[email protected]>
opensearch-project · Oct 22, 2024 · 3902aef · 3902aef
1 parent 48220eb
commit 3902aef
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -38,6 +38,7 @@ This section is for maintaining a changelog for all breaking changes for the cli
 ## [Unreleased 2.x]
 
 ### Added
+- Added `minScore` and `maxDistance` to `KnnQuery` ([#1166](https://github.com/opensearch-project/opensearch-java/pull/1166))
 
 ### Dependencies
 

diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQuery.java b/java-client/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQuery.java
@@ -23,7 +23,12 @@
 public class KnnQuery extends QueryBase implements QueryVariant {
     private final String field;
     private final float[] vector;
-    private final int k;
+    @Nullable
+    private final Integer k;
+    @Nullable
+    private final Float minScore;
+    @Nullable
+    private final Float maxDistance;
     @Nullable
     private final Query filter;
 
@@ -32,7 +37,9 @@ private KnnQuery(Builder builder) {
 
         this.field = ApiTypeHelper.requireNonNull(builder.field, this, "field");
         this.vector = ApiTypeHelper.requireNonNull(builder.vector, this, "vector");
-        this.k = ApiTypeHelper.requireNonNull(builder.k, this, "k");
+        this.k = builder.k;
+        this.minScore = builder.minScore;
+        this.maxDistance = builder.maxDistance;
         this.filter = builder.filter;
     }
 
@@ -66,13 +73,29 @@ public final float[] vector() {
     }
 
     /**
-     * Required - The number of neighbors the search of each graph will return.
+     * Optional - The number of neighbors the search of each graph will return.
      * @return The number of neighbors to return.
      */
-    public final int k() {
+    public final Integer k() {
         return this.k;
     }
 
+    /**
+     * Optional - The minimum score allowed for the returned search results.
+     * @return The minimum score allowed for the returned search results.
+     */
+    private final Float minScore() {
+        return this.minScore;
+    }
+
+    /**
+     * Optional - The maximum distance allowed between the vector and each of the returned search results.
+     * @return The maximum distance allowed between the vector and each ofthe returned search results.
+     */
+    private final Float maxDistance() {
+        return this.maxDistance;
+    }
+
     /**
      * Optional - A query to filter the results of the query.
      * @return The filter query.
@@ -97,7 +120,17 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
         }
         generator.writeEnd();
 
-        generator.write("k", this.k);
+        if (this.k != null) {
+            generator.write("k", this.k);
+        }
+
+        if (this.minScore != null) {
+            generator.write("min_score", this.minScore);
+        }
+
+        if (this.maxDistance != null) {
+            generator.write("max_distance", this.maxDistance);
+        }
 
         if (this.filter != null) {
             generator.writeKey("filter");
@@ -108,7 +141,7 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
     }
 
     public Builder toBuilder() {
-        return toBuilder(new Builder()).field(field).vector(vector).k(k).filter(filter);
+        return toBuilder(new Builder()).field(field).vector(vector).k(k).minScore(minScore).maxDistance(maxDistance).filter(filter);
     }
 
     /**
@@ -122,6 +155,10 @@ public static class Builder extends QueryBase.AbstractBuilder<Builder> implement
         @Nullable
         private Integer k;
         @Nullable
+        private Float minScore;
+        @Nullable
+        private Float maxDistance;
+        @Nullable
         private Query filter;
 
         /**
@@ -156,6 +193,28 @@ public Builder k(@Nullable Integer k) {
             return this;
         }
 
+        /**
+         * Optional - The minimum score allowed for the returned search results.
+         *
+         * @param minScore The minimum score allowed for the returned search results.
+         * @return This builder.
+         */
+        public Builder minScore(@Nullable Float minScore) {
+            this.minScore = minScore;
+            return this;
+        }
+
+        /**
+         * Optional - The maximum distance allowed between the vector and each of the returned search results.
+         *
+         * @param maxDistance The maximum distance allowed between the vector and each ofthe returned search results.
+         * @return This builder.
+         */
+        public Builder maxDistance(@Nullable Float maxDistance) {
+            this.maxDistance = maxDistance;
+            return this;
+        }
+
         /**
          * Optional - A query to filter the results of the knn query.
          *
@@ -201,6 +260,8 @@ protected static void setupKnnQueryDeserializer(ObjectDeserializer<Builder> op)
             b.vector(vector);
         }, JsonpDeserializer.arrayDeserializer(JsonpDeserializer.floatDeserializer()), "vector");
         op.add(Builder::k, JsonpDeserializer.integerDeserializer(), "k");
+        op.add(Builder::minScore, JsonpDeserializer.floatDeserializer(), "min_score");
+        op.add(Builder::maxDistance, JsonpDeserializer.floatDeserializer(), "max_distance");
         op.add(Builder::filter, Query._DESERIALIZER, "filter");
 
         op.setKey(Builder::field, JsonpDeserializer.stringDeserializer());

diff --git a/...-client/src/test/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQueryTest.java b/...-client/src/test/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQueryTest.java
@@ -14,7 +14,7 @@
 public class KnnQueryTest extends ModelTestCase {
     @Test
     public void toBuilder() {
-        KnnQuery origin = new KnnQuery.Builder().field("field").vector(new float[] { 1.0f }).k(1).build();
+        KnnQuery origin = new KnnQuery.Builder().field("field").vector(new float[] { 1.0f }).k(1).minScore(0.0f).maxDistance(1.0f).build();
         KnnQuery copied = origin.toBuilder().build();
 
         assertEquals(toJson(copied), toJson(origin));

diff --git a/java-client/src/test/java/org/opensearch/client/opensearch/model/VariantsTest.java b/java-client/src/test/java/org/opensearch/client/opensearch/model/VariantsTest.java
@@ -282,7 +282,7 @@ public void testHybridQuery() {
         assertEquals(100, searchRequest.query().hybrid().queries().get(1).neural().k());
         assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(2).knn().field());
         assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().vector().length);
-        assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().k());
+        assertEquals(Integer.valueOf(2), searchRequest.query().hybrid().queries().get(2).knn().k());
     }
 
     @Test
@@ -304,6 +304,6 @@ public void testHybridQueryFromJson() {
         assertEquals(100, searchRequest.query().hybrid().queries().get(1).neural().k());
         assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(2).knn().field());
         assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().vector().length);
-        assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().k());
+        assertEquals(Integer.valueOf(2), searchRequest.query().hybrid().queries().get(2).knn().k());
     }
 }