Tweaks: byte tests, JMH/benchmark polish, query javadoc

Author: iprithv

lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/DiversifyingChildrenFloatKnnJoinBenchmark.java

@@ -37,6 +37,7 @@
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery;
 import org.apache.lucene.search.join.QueryBitSetProducer;
+import org.apache.lucene.search.knn.KnnSearchStrategy;
 import org.apache.lucene.store.ByteBuffersDirectory;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.VectorUtil;
@@ -60,39 +61,38 @@
  * index (children + parent marker per block), using the default HNSW approximate path ({@code
  * childFilter == null}).
  *
- * <p>The {@code blockRescore} parameter switches the feature on/off so both modes can be compared
+ * <p>The {@code rescoreBlocks} parameter switches the feature on/off so both modes can be compared
  * in a single run (see <a href="https://github.com/apache/lucene/issues/15839">LUCENE-15839</a>).
  * Extra work scales roughly with {@code topK * childrenPerParent}.
  *
- * <p>Indicative results on Apple M-series, JDK 25, dim=96, topK=64, 4096 parent blocks (lower is
- * better):
+ * <p>Indicative results — 3 forks, 5 warmup / 10 measurement iterations, JDK 25, {@code -Xmx2g},
+ * dim=96, topK=64, 4096 parent blocks (lower is better):
  *
  * <pre>
- * blockRescore  childrenPerParent  Score (ms/op)
- * false         8                  0.123
- * false         32                 0.226
- * false         64                 0.254
- * true          8                  0.151  (+23%)
- * true          32                 0.316  (+40%)
- * true          64                 0.412  (+62%)
+ * childrenPerParent  rescoreBlocks=false    rescoreBlocks=true
+ * 8                  0.117 ± 0.002 ms/op    0.149 ± 0.002 ms/op  (+27%)
+ * 32                 0.237 ± 0.006 ms/op    0.326 ± 0.009 ms/op  (+38%)
+ * 64                 0.259 ± 0.005 ms/op    0.426 ± 0.013 ms/op  (+64%)
  * </pre>
  *
+ * <p>Overhead grows with block width (and with {@code topK}).
+ *
  * <p>Example:
  *
  * <pre>{@code
  * ./gradlew :lucene:benchmark-jmh:assemble
  * cd lucene/benchmark-jmh/build/benchmarks
  * java -jar lucene-benchmark-jmh-*-SNAPSHOT.jar DiversifyingChildrenFloatKnnJoin \\
- *     -f 2 -wi 3 -i 8 -tu ms
+ *     -f 3 -wi 5 -i 10 -tu ms
  * }</pre>
  */
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @State(Scope.Benchmark)
-@Warmup(iterations = 3, time = 1)
-@Measurement(iterations = 6, time = 1)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 10, time = 1)
 @Fork(
-    value = 2,
+    value = 3,
     jvmArgsAppend = {
       "-Xmx2g",
       "-Xms2g",
@@ -117,12 +117,12 @@ public class DiversifyingChildrenFloatKnnJoinBenchmark {
 
   /**
    * Whether to enable post-HNSW block rescoring. When {@code true}, after HNSW search all children
-   * in each found parent's block are scored to guarantee the best child is returned. Compare
-   * {@code false} (baseline / no rescoring) against {@code true} (rescoring enabled) to measure
-   * latency overhead.
+   * in each found parent's block are scored to guarantee the best child is returned. Compare {@code
+   * false} (baseline / no rescoring) against {@code true} (rescoring enabled) to measure latency
+   * overhead.
    */
   @Param({"false", "true"})
-  public boolean blockRescore;
+  public boolean rescoreBlocks;
 
   private Directory directory;
   private IndexSearcher searcher;
@@ -176,21 +176,18 @@ public void setupTrial() throws IOException {
     searcher = new IndexSearcher(reader);
     BitSetProducer parentsFilter =
         new QueryBitSetProducer(new TermQuery(new Term("docType", "_parent")));
+    // [1, 0, ..., 0] is already L2-normalized.
     float[] queryVector = new float[dimension];
     queryVector[0] = 1f;
-    for (int i = 1; i < dimension; i++) {
-      queryVector[i] = 0f;
-    }
-    VectorUtil.l2normalize(queryVector, false);
     diversifyingJoinQuery =
         new DiversifyingChildrenFloatKnnVectorQuery(
             "vec",
             queryVector,
             null,
             topK,
             parentsFilter,
-            org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT,
-            blockRescore);
+            KnnSearchStrategy.Hnsw.DEFAULT,
+            rescoreBlocks);
   }
 
   @TearDown(Level.Trial)

lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java

@@ -60,7 +60,7 @@ public class DiversifyingChildrenByteKnnVectorQuery extends KnnByteVectorQuery {
   private final Query childFilter;
   private final int k;
   private final byte[] query;
-  private final boolean blockRescore;
+  private final boolean rescoreBlocks;
 
   /**
    * Create a ToParentBlockJoinByteVectorQuery.
@@ -102,18 +102,23 @@ public DiversifyingChildrenByteKnnVectorQuery(
   /**
    * Create a DiversifyingChildrenByteKnnVectorQuery with optional post-HNSW block rescoring.
    *
-   * <p>When {@code blockRescore} is {@code true}, after the approximate HNSW search completes, all
-   * children in each found parent's block are scored to guarantee the truly best child is returned.
-   * See {@link DiversifyingChildrenFloatKnnVectorQuery#DiversifyingChildrenFloatKnnVectorQuery(
-   * String, float[], Query, int, BitSetProducer, KnnSearchStrategy, boolean)} for details.
+   * <p>When {@code rescoreBlocks} is {@code true}, after the approximate HNSW search completes, all
+   * children in each found parent's block are scored to guarantee the truly best child is returned
+   * — not merely the sibling the graph traversal happened to reach first. This adds O(k &times;
+   * childrenPerParent) extra scoring work; enable it when block sizes are small or result quality
+   * is more important than latency.
+   *
+   * <p>This applies only to the approximate (HNSW) search path. When the index is small enough that
+   * Lucene falls back to exact search, all children are already scored exhaustively and no
+   * additional rescoring is performed.
    *
    * @param field the query field
    * @param query the vector query
    * @param childFilter the child filter
    * @param k how many parent documents to return given the matching children
    * @param parentsFilter Filter identifying the parent documents.
    * @param searchStrategy the search strategy to use.
-   * @param blockRescore if {@code true}, enables post-HNSW block rescoring.
+   * @param rescoreBlocks if {@code true}, enables post-HNSW block rescoring.
    * @lucene.experimental
    */
   public DiversifyingChildrenByteKnnVectorQuery(
@@ -123,13 +128,13 @@ public DiversifyingChildrenByteKnnVectorQuery(
       int k,
       BitSetProducer parentsFilter,
       KnnSearchStrategy searchStrategy,
-      boolean blockRescore) {
+      boolean rescoreBlocks) {
     super(field, query, k, childFilter, searchStrategy);
     this.childFilter = childFilter;
     this.parentsFilter = parentsFilter;
     this.k = k;
     this.query = query;
-    this.blockRescore = blockRescore;
+    this.rescoreBlocks = rescoreBlocks;
   }
 
   @Override
@@ -204,7 +209,7 @@ protected TopDocs approximateSearch(
     }
     context.reader().searchNearestVectors(field, query, collector, acceptDocs);
     TopDocs results = collector.topDocs();
-    if (!blockRescore || results.scoreDocs.length == 0) {
+    if (!rescoreBlocks || results.scoreDocs.length == 0) {
       return results;
     }
     BitSet parentBitSet = parentsFilter.getBitSet(context);
@@ -243,15 +248,15 @@ public boolean equals(Object o) {
     if (!super.equals(o)) return false;
     DiversifyingChildrenByteKnnVectorQuery that = (DiversifyingChildrenByteKnnVectorQuery) o;
     return k == that.k
-        && blockRescore == that.blockRescore
+        && rescoreBlocks == that.rescoreBlocks
         && Objects.equals(parentsFilter, that.parentsFilter)
         && Objects.equals(childFilter, that.childFilter)
         && Arrays.equals(query, that.query);
   }
 
   @Override
   public int hashCode() {
-    int result = Objects.hash(super.hashCode(), parentsFilter, childFilter, k, blockRescore);
+    int result = Objects.hash(super.hashCode(), parentsFilter, childFilter, k, rescoreBlocks);
     result = 31 * result + Arrays.hashCode(query);
     return result;
   }

lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java

@@ -20,7 +20,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Comparator;
 import java.util.Objects;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.LeafReaderContext;
@@ -62,7 +61,7 @@ public class DiversifyingChildrenFloatKnnVectorQuery extends KnnFloatVectorQuery
   private final Query childFilter;
   private final int k;
   private final float[] query;
-  private final boolean blockRescore;
+  private final boolean rescoreBlocks;
 
   /**
    * Create a DiversifyingChildrenFloatKnnVectorQuery.
@@ -104,19 +103,23 @@ public DiversifyingChildrenFloatKnnVectorQuery(
   /**
    * Create a DiversifyingChildrenFloatKnnVectorQuery with optional post-HNSW block rescoring.
    *
-   * <p>When {@code blockRescore} is {@code true}, after the approximate HNSW search completes, all
-   * children in each found parent's block are scored to guarantee the truly best child is returned —
-   * not merely the sibling the graph traversal happened to reach first. This adds O(k &times;
-   * childrenPerParent) extra scoring work; enable it when block sizes are small or result quality is
-   * more important than latency.
+   * <p>When {@code rescoreBlocks} is {@code true}, after the approximate HNSW search completes, all
+   * children in each found parent's block are scored to guarantee the truly best child is returned
+   * — not merely the sibling the graph traversal happened to reach first. This adds O(k &times;
+   * childrenPerParent) extra scoring work; enable it when block sizes are small or result quality
+   * is more important than latency.
+   *
+   * <p>This applies only to the approximate (HNSW) search path. When the index is small enough that
+   * Lucene falls back to exact search, all children are already scored exhaustively and no
+   * additional rescoring is performed.
    *
    * @param field the query field
    * @param query the vector query
    * @param childFilter the child filter
    * @param k how many parent documents to return given the matching children
    * @param parentsFilter Filter identifying the parent documents.
    * @param searchStrategy the search strategy to use.
-   * @param blockRescore if {@code true}, enables post-HNSW block rescoring.
+   * @param rescoreBlocks if {@code true}, enables post-HNSW block rescoring.
    * @lucene.experimental
    */
   public DiversifyingChildrenFloatKnnVectorQuery(
@@ -126,13 +129,13 @@ public DiversifyingChildrenFloatKnnVectorQuery(
       int k,
       BitSetProducer parentsFilter,
       KnnSearchStrategy searchStrategy,
-      boolean blockRescore) {
+      boolean rescoreBlocks) {
     super(field, query, k, childFilter, searchStrategy);
     this.childFilter = childFilter;
     this.parentsFilter = parentsFilter;
     this.k = k;
     this.query = query;
-    this.blockRescore = blockRescore;
+    this.rescoreBlocks = rescoreBlocks;
   }
 
   @Override
@@ -206,7 +209,7 @@ protected TopDocs approximateSearch(
     }
     context.reader().searchNearestVectors(field, query, collector, acceptDocs);
     TopDocs results = collector.topDocs();
-    if (!blockRescore || results.scoreDocs.length == 0) {
+    if (!rescoreBlocks || results.scoreDocs.length == 0) {
       return results;
     }
     BitSet parentBitSet = parentsFilter.getBitSet(context);
@@ -226,13 +229,10 @@ protected TopDocs approximateSearch(
 
   /**
    * For each parent already found by approximate search, scores all children in that parent's block
-   * to ensure the truly best child is returned — not merely the sibling the graph traversal happened
-   * to reach first. Children are processed in ascending docId order so the sequential {@link
-   * VectorScorer} only advances forward. Extra nodes scored are added to {@link
+   * to ensure the truly best child is returned — not merely the sibling the graph traversal
+   * happened to reach first. Children are processed in ascending docId order so the sequential
+   * {@link VectorScorer} only advances forward. Extra nodes scored are added to {@link
    * TotalHits#value()}.
-   *
-   * <p>This method is package-private so that {@link DiversifyingChildrenByteKnnVectorQuery} can
-   * reuse the same implementation rather than duplicating it.
    */
   static TopDocs blockRescore(
       TopDocs results, AcceptDocs acceptDocs, BitSet parentBitSet, VectorScorer scorer)
@@ -243,7 +243,7 @@ static TopDocs blockRescore(
     // Sort by docId so parent blocks are visited in ascending order — the forward-only
     // VectorScorer cannot go backwards.
     ScoreDoc[] scoreDocs = results.scoreDocs.clone();
-    Arrays.sort(scoreDocs, Comparator.comparingInt(sd -> sd.doc));
+    Arrays.sort(scoreDocs, (a, b) -> Integer.compare(a.doc, b.doc));
 
     long extraVisited = 0;
     for (ScoreDoc scoreDoc : scoreDocs) {
@@ -256,10 +256,11 @@ static TopDocs blockRescore(
           continue;
         }
         if (scorerIter.advance(child) == child) {
-          // Don't double-count the child HNSW already visited.
-          if (child != hnswBestChild) {
-            extraVisited++;
+          if (child == hnswBestChild) {
+            // Advance past the child HNSW already scored; no need to re-compute.
+            continue;
           }
+          extraVisited++;
           float s = scorer.score();
           if (s > scoreDoc.score) {
             scoreDoc.score = s;
@@ -293,15 +294,15 @@ public boolean equals(Object o) {
     if (!super.equals(o)) return false;
     DiversifyingChildrenFloatKnnVectorQuery that = (DiversifyingChildrenFloatKnnVectorQuery) o;
     return k == that.k
-        && blockRescore == that.blockRescore
+        && rescoreBlocks == that.rescoreBlocks
         && Objects.equals(parentsFilter, that.parentsFilter)
         && Objects.equals(childFilter, that.childFilter)
         && Arrays.equals(query, that.query);
   }
 
   @Override
   public int hashCode() {
-    int result = Objects.hash(super.hashCode(), parentsFilter, childFilter, k, blockRescore);
+    int result = Objects.hash(super.hashCode(), parentsFilter, childFilter, k, rescoreBlocks);
     result = 31 * result + Arrays.hashCode(query);
     return result;
   }