Feature/collaborative hnsw search

lucene/CHANGES.txt

@@ -56,11 +56,14 @@ API Changes
 
 * GITHUB#14615 : Remove unnecessary public methods in FuzzySet (Greg Miller)
 
-* GITHUB#15295 : Switched to a fixed CFS threshold (Shubham Sharma)
+* GITHUB#15295 : Switched to a fixed CFS threshold(Shubham Sharma)
 
 New Features
 ---------------------
 
+* GITHUB#KNN-COLLAB: Introduce Collaborative HNSW search, allowing dynamic threshold
+  updates from collectors to enable cluster-wide search pruning. (ai-pipestream)
+
 * GITHUB#15505: Upgrade snowball to 2d2e312df56f2ede014a4ffb3e91e6dea43c24be. New stemmer: PolishStemmer (and
   PolishSnowballAnalyzer in the stempel package) (Justas Sakalauskas, Dawid Weiss)
 
@@ -1774,7 +1777,7 @@ Optimizations
 
 * GITHUB#13184: Make the HitQueue size more appropriate for KNN exact search (Pan Guixin)
 
-* GITHUB#13199: Speed up dynamic pruning by breaking point estimation when threshold get exceeded. (Guo Feng)
+* GITHUB#13199: Speed up dynamic pruning by breaking point estimation when thresholdget exceeded. (Guo Feng)
 
 * GITHUB#13203: Speed up writeGroupVInts (Zhang Chao)
 

lucene/core/src/java/org/apache/lucene/search/CollaborativeKnnCollector.java

@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search;
+
+import java.util.concurrent.atomic.LongAccumulator;
+import java.util.function.IntUnaryOperator;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
+
+/**
+ * A {@link KnnCollector} that allows for collaborative search. PRUNING BASED ON GLOBAL FLOOR vs
+ * LOCAL MAX.
+ */
+public class CollaborativeKnnCollector extends KnnCollector.Decorator {
+
+  private static final IntUnaryOperator IDENTITY_MAPPER = docId -> docId;
+  private static final int GLOBAL_BAR_MIN_VISITS = 100;
+  private static final float GLOBAL_BAR_TERMINATION_SLACK = 0.0001f;
+
+  private final LongAccumulator minScoreAcc;
+  private final int docBase;
+  private final IntUnaryOperator docIdMapper;
+
+  private float localMaxScore = Float.NEGATIVE_INFINITY;
+  private float lastSharedScore = Float.NEGATIVE_INFINITY;
+
+  /** Convenience constructor for tests. */
+  public CollaborativeKnnCollector(
+      int k, int visitLimit, LongAccumulator minScoreAcc, int docBase) {
+    this(new TopKnnCollector(k, visitLimit), minScoreAcc, docBase, IDENTITY_MAPPER);
+  }
+
+  public CollaborativeKnnCollector(
+      int k,
+      int visitLimit,
+      LongAccumulator minScoreAcc,
+      int docBase,
+      IntUnaryOperator docIdMapper) {
+    this(new TopKnnCollector(k, visitLimit), minScoreAcc, docBase, docIdMapper);
+  }
+
+  public CollaborativeKnnCollector(
+      int k,
+      int visitLimit,
+      KnnSearchStrategy searchStrategy,
+      LongAccumulator minScoreAcc,
+      int docBase,
+      IntUnaryOperator docIdMapper) {
+    this(new TopKnnCollector(k, visitLimit, searchStrategy), minScoreAcc, docBase, docIdMapper);
+  }
+
+  private CollaborativeKnnCollector(
+      KnnCollector delegate,
+      LongAccumulator minScoreAcc,
+      int docBase,
+      IntUnaryOperator docIdMapper) {
+    super(delegate);
+    this.minScoreAcc = minScoreAcc;
+    this.docBase = docBase;
+    this.docIdMapper = docIdMapper;
+  }
+
+  @Override
+  public float minCompetitiveSimilarity() {
+    // Pathfinding always uses local bar
+    return super.minCompetitiveSimilarity();
+  }
+
+  @Override
+  public boolean earlyTerminated() {
+    if (super.earlyTerminated()) return true;
+    if (visitedCount() < GLOBAL_BAR_MIN_VISITS) return false;
+
+    long globalFloorCode = minScoreAcc.get();
+    if (globalFloorCode == Long.MIN_VALUE) return false;
+
+    float globalFloorScore = DocScoreEncoder.toScore(globalFloorCode);
+
+    // CRITICAL: Only stop if our BEST hit is worse than the global floor.
+    // If localMax < globalFloor, it's impossible for this shard to make the Top K.
+    return localMaxScore > Float.NEGATIVE_INFINITY
+        && localMaxScore < (globalFloorScore - GLOBAL_BAR_TERMINATION_SLACK);
+  }
+
+  @Override
+  public boolean collect(int docId, float similarity) {
+    boolean collected = super.collect(docId, similarity);
+
+    if (similarity > localMaxScore) {
+      localMaxScore = similarity;
+    }
+
+    if (collected) {
+      float floorScore = super.minCompetitiveSimilarity();
+      if (floorScore > Float.NEGATIVE_INFINITY && floorScore > lastSharedScore + 0.0001f) {
+
+        int absoluteDocId = docId + docBase;
+        minScoreAcc.accumulate(
+            DocScoreEncoder.encode(docIdMapper.applyAsInt(absoluteDocId), floorScore));
+        lastSharedScore = floorScore;
+      }
+    }
+    return collected;
+  }
+
+  public static float toScore(long value) {
+    return DocScoreEncoder.toScore(value);
+  }
+
+  public static long encode(int docId, float score) {
+    return DocScoreEncoder.encode(docId, score);
+  }
+}

lucene/core/src/java/org/apache/lucene/search/knn/CollaborativeKnnCollectorManager.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.search.knn;
+
+import java.io.IOException;
+import java.util.concurrent.atomic.LongAccumulator;
+import java.util.function.IntUnaryOperator;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.CollaborativeKnnCollector;
+import org.apache.lucene.search.KnnCollector;
+
+/**
+ * A {@link KnnCollectorManager} that creates {@link CollaborativeKnnCollector} instances sharing a
+ * single {@link LongAccumulator} for global pruning across segments, gated by topological hints.
+ *
+ * @lucene.experimental
+ */
+public class CollaborativeKnnCollectorManager implements KnnCollectorManager {
+
+  private final int k;
+  private final LongAccumulator minScoreAcc;
+  private final IntUnaryOperator docIdMapper;
+
+  /**
+   * Create a new CollaborativeKnnCollectorManager
+   *
+   * @param k number of neighbors to collect
+   * @param minScoreAcc shared accumulator for global pruning
+   */
+  public CollaborativeKnnCollectorManager(int k, LongAccumulator minScoreAcc) {
+    this(k, minScoreAcc, docId -> docId);
+  }
+
+  /** Create a new CollaborativeKnnCollectorManager with a docId mapper */
+  public CollaborativeKnnCollectorManager(
+      int k, LongAccumulator minScoreAcc, IntUnaryOperator docIdMapper) {
+    this.k = k;
+    this.minScoreAcc = minScoreAcc;
+    this.docIdMapper = docIdMapper;
+  }
+
+  @Override
+  public KnnCollector newCollector(
+      int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
+      throws IOException {
+    return new CollaborativeKnnCollector(
+        k, visitedLimit, searchStrategy, minScoreAcc, context.docBase, docIdMapper);
+  }
+}

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java

@@ -305,6 +305,13 @@ void searchLevel(
     // We should allow exploring equivalent minAcceptedSimilarity values at least once
     boolean shouldExploreMinSim = true;
     while (candidates.size() > 0 && results.earlyTerminated() == false) {
+      // Update the threshold dynamically from the collector to allow external pruning.
+      float liveMinSimilarity = results.minCompetitiveSimilarity();
+      if (liveMinSimilarity > minAcceptedSimilarity) {
+        minAcceptedSimilarity = liveMinSimilarity;
+        shouldExploreMinSim = true;
+      }
+
       // get the best candidate (closest or best scoring)
       float topCandidateSimilarity = candidates.topScore();
       if (topCandidateSimilarity < minAcceptedSimilarity) {