add distanceThreshold option for nearestNeighbor operator

1 files changed, 10 insertions, 0 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/query/NearestNeighborItem.java b/container-search/src/main/java/com/yahoo/prelude/query/NearestNeighborItem.java
index e237463582f..4a827813e81 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/NearestNeighborItem.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/NearestNeighborItem.java
@@ -22,6 +22,7 @@ public class NearestNeighborItem extends SimpleTaggableItem {
 
     private int targetNumHits = 0;
     private int hnswExploreAdditionalHits = 0;
+    private double distanceThreshold = Double.POSITIVE_INFINITY;
     private boolean approximate = true;
     private String field;
     private final String queryTensorName;
@@ -37,6 +38,9 @@ public class NearestNeighborItem extends SimpleTaggableItem {
     /** Returns the field name */
     public String getIndexName() { return field; }
 
+    /** Returns the distance threshold for nearest-neighbor hits */
+    public double getDistanceThreshold () { return this.distanceThreshold ; }
+
     /** Returns the number of extra hits to explore in HNSW algorithm */
     public int getHnswExploreAdditionalHits() { return hnswExploreAdditionalHits; }
 
@@ -49,6 +53,9 @@ public class NearestNeighborItem extends SimpleTaggableItem {
     /** Set the K number of hits to produce */
     public void setTargetNumHits(int target) { this.targetNumHits = target; }
 
+    /** Set the distance threshold for nearest-neighbor hits */
+    public void setDistanceThreshold(double threshold) { this.distanceThreshold = threshold; }
+
     /** Set the number of extra hits to explore in HNSW algorithm */
     public void setHnswExploreAdditionalHits(int num) { this.hnswExploreAdditionalHits = num; }
 
@@ -75,6 +82,7 @@ public class NearestNeighborItem extends SimpleTaggableItem {
         IntegerCompressor.putCompressedPositiveNumber(targetNumHits, buffer);
         IntegerCompressor.putCompressedPositiveNumber((approximate ? 1 : 0), buffer);
         IntegerCompressor.putCompressedPositiveNumber(hnswExploreAdditionalHits, buffer);
+        buffer.putDouble(distanceThreshold);
         return 1;  // number of encoded stack dump items
     }
 
@@ -83,6 +91,7 @@ public class NearestNeighborItem extends SimpleTaggableItem {
         buffer.append("{field=").append(field);
         buffer.append(",queryTensorName=").append(queryTensorName);
         buffer.append(",hnsw.exploreAdditionalHits=").append(hnswExploreAdditionalHits);
+        buffer.append(",distanceThreshold=").append(distanceThreshold);
         buffer.append(",approximate=").append(approximate);
         buffer.append(",targetHits=").append(targetNumHits).append("}");
     }
@@ -93,6 +102,7 @@ public class NearestNeighborItem extends SimpleTaggableItem {
         discloser.addProperty("field", field);
         discloser.addProperty("queryTensorName", queryTensorName);
         discloser.addProperty("hnsw.exploreAdditionalHits", hnswExploreAdditionalHits);
+        discloser.addProperty("distanceThreshold", distanceThreshold);
         discloser.addProperty("approximate", approximate);
         discloser.addProperty("targetHits", targetNumHits);
     }