resolve comments

src/main/java/com/yelp/nrtsearch/server/luceneserver/highlights/TopBoostOnlyFragmentsBuilderAdaptor.java

@@ -15,6 +15,7 @@
  */
 package com.yelp.nrtsearch.server.luceneserver.highlights;
 
+import java.util.ArrayList;
 import java.util.List;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.search.highlight.Encoder;
@@ -24,6 +25,13 @@
 import org.apache.lucene.search.vectorhighlight.FieldFragList.WeightedFragInfo.SubInfo;
 import org.apache.lucene.search.vectorhighlight.FieldPhraseList.WeightedPhraseInfo.Toffs;
 
+/**
+ * Adapter for {@link org.apache.lucene.search.vectorhighlight.FragmentsBuilder} that highlights
+ * only the top matched phrases based on the boost value in the query. This adapter does not alter
+ * the order or score of the generated fragments. All phrases contribute to scoring if the
+ * innerBaseFragmentsBuilder is a {@link
+ * org.apache.lucene.search.vectorhighlight.ScoreOrderFragmentsBuilder}.
+ */
 public class TopBoostOnlyFragmentsBuilderAdaptor extends BaseFragmentsBuilder {
   private final BaseFragmentsBuilder innerBaseFragmentsBuilder;
   private final boolean topBoostOnly;
@@ -38,11 +46,16 @@ public TopBoostOnlyFragmentsBuilderAdaptor(
     this.topBoostOnly = topBoostOnly;
   }
 
+  /** Delegates the inner FragmentsBuilder to determine the fragment order. */
   @Override
   public List<WeightedFragInfo> getWeightedFragInfoList(List<WeightedFragInfo> src) {
     return innerBaseFragmentsBuilder.getWeightedFragInfoList(src);
   }
 
+  /**
+   * Creates a fragment containing only the top boost phrase if the `topBoostOnly` flag is set.
+   * Otherwise, it delegates to the base implementation.
+   */
   @Override
   protected String makeFragment(
       StringBuilder buffer,
@@ -62,9 +75,21 @@ protected String makeFragment(
         getFragmentSourceMSO(
             buffer, index, values, s, fragInfo.getEndOffset(), modifiedStartOffset);
     int srcIndex = 0;
-    double topBoostValue =
-        fragInfo.getSubInfos().stream().map(SubInfo::getBoost).max(Float::compare).orElse(0f);
+    // filter out the phrases with lower boost at the fragment creation time only
+    float topBoostValue = 0;
+    List<SubInfo> topSubInfoList = new ArrayList<>();
     for (SubInfo subInfo : fragInfo.getSubInfos()) {
+      float boost = subInfo.getBoost();
+      if (boost > topBoostValue) {
+        topBoostValue = boost;
+        topSubInfoList.clear();
+        topSubInfoList.add(subInfo);
+      } else if (boost == topBoostValue) {
+        topSubInfoList.add(subInfo);
+      }
+    }
+
+    for (SubInfo subInfo : topSubInfoList) {
       if (subInfo.getBoost() < topBoostValue) {
         continue;
       }