avg aggregator fix, iterators refactoring

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · Oct 1, 2024 · 89de2f8 · 89de2f8
1 parent 2ac5915
commit 89de2f8
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -211,7 +211,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
             // Iterate over the FixedBitSet
             for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
                 // Advance to the entryId in the valuesIterator
-                if (!valuesIterator.advanceExact(bit)) {
+                if (valuesIterator.advanceExact(bit) == false) {
                     continue;  // Skip if no more entries
                 }
 

diff --git a/...nsearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/...nsearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
@@ -21,7 +21,7 @@
  * @opensearch.experimental
  */
 @ExperimentalApi
-public class StarTreeValuesIterator {
+public abstract class StarTreeValuesIterator {
 
     public static final int NO_MORE_ENTRIES = Integer.MAX_VALUE;
     protected final DocIdSetIterator docIdSetIterator;

diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -46,7 +46,6 @@
 import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
 import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
 import org.opensearch.index.fielddata.SortedNumericDoubleValues;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.aggregations.Aggregator;
@@ -177,8 +176,7 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
             // Iterate over the FixedBitSet
             for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
                 // Advance to the bit (entryId) in the valuesIterator
-                if (sumValuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES
-                    || countValueIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+                if ((sumValuesIterator.advanceExact(bit) && countValueIterator.advanceExact(bit)) == false) {
                     continue;  // Skip if no more entries
                 }
 

diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -62,7 +62,9 @@ public FixedBitSet getStarTreeResult() throws IOException {
 
         // Initialize FixedBitSet with size maxMatchedDoc + 1
         FixedBitSet bitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
-        StarTreeValuesIterator starTreeValuesIterator = new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator());
+        SortedNumericStarTreeValuesIterator starTreeValuesIterator = new SortedNumericStarTreeValuesIterator(
+            starTreeResult._matchedDocIds.build().iterator()
+        );
 
         // No matches, return an empty FixedBitSet
         if (starTreeResult.maxMatchedDoc == -1) {