Using advanceExact instead of advance due to AssertionError at IndexedDISI

expani · expani · commit 8f9477a425dc · 2025-01-25T16:53:57.000-08:00
Signed-off-by: expani &lt;anijainc@amazon.com&gt;
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java
@@ -34,6 +34,11 @@ public long value() throws IOException {
         return nextOrd();
     }
 
+    @Override
+    public boolean advanceExact(int target) throws IOException {
+        return ((SortedSetDocValues) docIdSetIterator).advanceExact(target);
+    }
+
     public long nextOrd() throws IOException {
         return ((SortedSetDocValues) docIdSetIterator).nextOrd();
     }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
@@ -48,4 +48,6 @@ public long cost() {
 
     public abstract long value() throws IOException;
 
+    public abstract boolean advanceExact(int target) throws IOException;
+
 }
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java b/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java
@@ -18,7 +18,6 @@
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator;
 import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.filter.DimensionFilter;
@@ -58,8 +57,6 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, StarT
 
         StarTreeResult starTreeResult = traverseStarTree(starTreeValues, starTreeFilter);
 
-        //iterateOverAllDimensionValues(starTreeValues);
-
         // Initialize FixedBitSet with size maxMatchedDoc + 1
         FixedBitSet bitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
         SortedNumericStarTreeValuesIterator starTreeValuesIterator = new SortedNumericStarTreeValuesIterator(
@@ -95,7 +92,7 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, StarT
                 for (int entryId = bitSet.nextSetBit(0); entryId != DocIdSetIterator.NO_MORE_DOCS; entryId = (entryId + 1 < bitSet.length())
                     ? bitSet.nextSetBit(entryId + 1)
                     : DocIdSetIterator.NO_MORE_DOCS) {
-                    if (valuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+                    if (valuesIterator.advanceExact(entryId)) {
                         long value = valuesIterator.value();
                         for (DimensionFilter dimensionFilter : dimensionFilters) {
                             if (dimensionFilter.matchDimValue(value, starTreeValues)) {
@@ -114,31 +111,6 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, StarT
         return bitSet;  // Return the final FixedBitSet with all matches
     }
 
-    private static void iterateOverAllDimensionValues(StarTreeValues starTreeValues) throws IOException {
-        for (String dimensionName : starTreeValues.getStarTreeField().getDimensionNames()) {
-            StarTreeValuesIterator starTreeValuesIterator = starTreeValues.getDimensionValuesIterator(dimensionName);
-            int docId = 0;
-            List<Integer> docIds = new ArrayList<>();
-            List<Object> values = new ArrayList<>();
-            while (starTreeValuesIterator.advance(docId) != NO_MORE_DOCS) {
-                docIds.add(starTreeValuesIterator.entryId());
-                        if (starTreeValuesIterator instanceof SortedNumericStarTreeValuesIterator) {
-                            values.add(((SortedNumericStarTreeValuesIterator) starTreeValuesIterator).nextValue());
-                        } else if (starTreeValuesIterator instanceof SortedSetStarTreeValuesIterator) {
-                            values.add(((SortedSetStarTreeValuesIterator) starTreeValuesIterator).nextOrd());
-                        }
-                docId++;
-            }
-            System.out.printf(
-                "\n\nDimension Values for dimension %s has size %s and are %s and docIds are %s",
-                dimensionName,
-                values.size(),
-                values,
-                docIds
-            );
-        }
-    }
-
     private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, StarTreeFilter starTreeFilter) throws IOException {
         DocIdSetBuilder docsWithField = new DocIdSetBuilder(starTreeValues.getStarTreeDocumentCount());
         DocIdSetBuilder.BulkAdder adder;
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -111,6 +111,7 @@ protected Codec getCodec() {
         final Logger testLogger = LogManager.getLogger(MetricAggregatorTests.class);
         MapperService mapperService;
         try {
+            // FIXME: Test for different max_leaf_docs with at least 1, 2-100, 101-10_000 for all query types
             mapperService = StarTreeDocValuesFormatTests.createMapperService(
                 StarTreeQueryTests.getExpandedMapping(randomIntBetween(1, 10_000), false)
             );
@@ -124,6 +125,7 @@ protected Codec getCodec() {
     public void testStarTreeDocValues() throws IOException {
         Directory directory = newDirectory();
         IndexWriterConfig conf = newIndexWriterConfig(null);
+        // FIXME: Test for different max_leaf_docs with at least 1, 2-100, 101-10_000 for all query types
         conf.setCodec(getCodec());
         conf.setMergePolicy(newLogMergePolicy());
         RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
@@ -139,6 +141,7 @@ public void testStarTreeDocValues() throws IOException {
         // Index 100 random documents
         for (int i = 0; i < totalDocs; i++) {
             Document doc = new Document();
+            // FIXME: Reduce the frequency of nulls to be with at least some non-null docs like after every 1-2 ?
             if (random.nextBoolean()) {
                 val = random.nextInt(10) - 5; // Random long between -5 and 4
                 doc.add(new LongField(SNDV, val, Field.Store.YES));
@@ -174,6 +177,7 @@ public void testStarTreeDocValues() throws IOException {
 
         MapperService mapperService = mapperServiceMock();
         CircuitBreakerService circuitBreakerService = new NoneCircuitBreakerService();
+        // FIXME : Register for other dimensions. Move to a common place where per dimension actions are done ?
         when(mapperService.fieldType(SNDV)).thenReturn(new NumberFieldMapper.NumberFieldType(SNDV, NumberFieldMapper.NumberType.LONG));
         when(mapperService.fieldType(DV)).thenReturn(new NumberFieldMapper.NumberFieldType(DV, NumberFieldMapper.NumberType.LONG));
         when(mapperService.fieldType(KEYWORD)).thenReturn(new KeywordFieldMapper.KeywordFieldType(KEYWORD));
@@ -184,6 +188,7 @@ public void testStarTreeDocValues() throws IOException {
             circuitBreakerService,
             new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), circuitBreakerService).withCircuitBreaking()
         );
+        // FIXME : Register for other dimensions. Move to a common place where per dimension actions are done ?
         when(queryShardContext.fieldMapper(SNDV)).thenReturn(
             new NumberFieldMapper.NumberFieldType(SNDV, NumberFieldMapper.NumberType.LONG)
         );
@@ -199,17 +204,19 @@ public void testStarTreeDocValues() throws IOException {
         ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
         AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
 
+        // FIXME : Add Float, Half Float and Double. Integer/Long whichever is missing
         List<Dimension> supportedDimensions = new LinkedList<>();
         supportedDimensions.add(new NumericDimension(SNDV));
         supportedDimensions.add(new NumericDimension(DV));
         supportedDimensions.add(new OrdinalDimension(KEYWORD));
 
         Query query = new MatchAllDocsQuery();
         // match-all query
+        // FIXME : Add support for numeric and keyword terms query
         QueryBuilder numericTermQueryBuilder = null; // no predicates
         QueryBuilder numericRangeQueryBuilder = null;
-        QueryBuilder keywordTermQueryBuilder = null;
-        QueryBuilder keywordRangeQueryBuilder = null;
+        QueryBuilder keywordTermQueryBuilder;
+        QueryBuilder keywordRangeQueryBuilder;
         testCase(
             indexSearcher,
             query,
@@ -455,10 +462,10 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
         QueryBuilder queryBuilder,
         T aggBuilder,
         CompositeIndexFieldInfo starTree,
-        List<Dimension> supportedDimensions,
-        List<Metric> supportedMetrics,
+        List<Dimension> supportedDimensions, // FIXME : Merge with the same input that goes to generating the codec.
+        List<Metric> supportedMetrics, // FIXME : Merge with the same input that goes to generating the codec.
         BiConsumer<V, V> verify,
-        AggregatorFactory aggregatorFactory,
+        AggregatorFactory aggregatorFactory, // TODO : Recheck if this can be done elsewhere.
         boolean assertCollectorEarlyTermination
     ) throws IOException {
         V starTreeAggregation = searchAndReduceStarTree(
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryTests.java
@@ -357,6 +357,8 @@ public static XContentBuilder getExpandedMapping(int maxLeafDocs, boolean skipSt
                 b.value("sdv");
                 b.endArray();
             }
+            // FIXME : Change to take dimension order and other inputs as method params.
+            // FIXME : Create default constants for the existing so other can call easily.
             b.startArray("ordered_dimensions");
             b.startObject();
             b.field("name", "sndv");
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -431,6 +431,7 @@ protected SearchContext createSearchContextWithStarTreeContext(
         when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes);
         when(searchContext.mapperService()).thenReturn(mapperService);
 
+        // FIXME : Directly supply the dimension as method param, we are resetting this again ?
         for (Dimension dimension : supportedDimensions) {
             if (dimension instanceof OrdinalDimension) {
                 MappedFieldType mappedFieldType = new KeywordFieldMapper.KeywordFieldType(dimension.getField(), Lucene.WHITESPACE_ANALYZER);

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,11 @@ public long value() throws IOException {`
`34`	`34`	`return nextOrd();`
`35`	`35`	`}`
`36`	`36`
	`37`	`+ @Override`
	`38`	`+ public boolean advanceExact(int target) throws IOException {`
	`39`	`+ return ((SortedSetDocValues) docIdSetIterator).advanceExact(target);`
	`40`	`+ }`
	`41`	`+`
`37`	`42`	`public long nextOrd() throws IOException {`
`38`	`43`	`return ((SortedSetDocValues) docIdSetIterator).nextOrd();`
`39`	`44`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,4 +48,6 @@ public long cost() {`
`48`	`48`
`49`	`49`	`public abstract long value() throws IOException;`
`50`	`50`
	`51`	`+ public abstract boolean advanceExact(int target) throws IOException;`
	`52`	`+`
`51`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -357,6 +357,8 @@ public static XContentBuilder getExpandedMapping(int maxLeafDocs, boolean skipSt`
`357`	`357`	`b.value("sdv");`
`358`	`358`	`b.endArray();`
`359`	`359`	`}`
	`360`	`+ // FIXME : Change to take dimension order and other inputs as method params.`
	`361`	`+ // FIXME : Create default constants for the existing so other can call easily.`
`360`	`362`	`b.startArray("ordered_dimensions");`
`361`	`363`	`b.startObject();`
`362`	`364`	`b.field("name", "sndv");`