interfaces fixes for star tree search

Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>
opensearch-project · Sep 3, 2024 · 7a1840a · 7a1840a
1 parent a5590d1
commit 7a1840a
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 14 deletions.
diff --git a/.../opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java b/.../opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java
@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.store.IndexInput;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.compositeindex.CompositeIndexMetadata;
 import org.opensearch.index.compositeindex.datacube.Metric;
 import org.opensearch.index.compositeindex.datacube.MetricStat;
@@ -30,6 +31,7 @@
  *
  * @opensearch.experimental
  */
+@ExperimentalApi
 public class StarTreeMetadata extends CompositeIndexMetadata {
     private static final Logger logger = LogManager.getLogger(StarTreeMetadata.class);
 

diff --git a/...arch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java b/...arch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java
@@ -201,7 +201,6 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOExce
         StarTreeNode resultStarTreeNode = null;
         if (null != dimensionValue) {
             resultStarTreeNode = binarySearchChild(dimensionValue);
-            assert null != resultStarTreeNode;
         }
         return resultStarTreeNode;
     }

diff --git a/...main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java b/...main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java
@@ -74,6 +74,11 @@ public class StarTreeValues implements CompositeIndexValues {
      */
     private final Map<String, String> attributes;
 
+    /**
+     * A metadata for the star-tree
+     */
+    private final StarTreeMetadata starTreeMetadata;
+
     /**
      * Constructs a new StarTreeValues object with the provided parameters.
      * Used for testing.
@@ -89,13 +94,15 @@ public StarTreeValues(
         StarTreeNode root,
         Map<String, Supplier<DocIdSetIterator>> dimensionDocValuesIteratorMap,
         Map<String, Supplier<DocIdSetIterator>> metricDocValuesIteratorMap,
-        Map<String, String> attributes
+        Map<String, String> attributes,
+        StarTreeMetadata compositeIndexMetadata
     ) {
         this.starTreeField = starTreeField;
         this.root = root;
         this.dimensionDocValuesIteratorMap = dimensionDocValuesIteratorMap;
         this.metricDocValuesIteratorMap = metricDocValuesIteratorMap;
         this.attributes = attributes;
+        this.starTreeMetadata = compositeIndexMetadata;
     }
 
     /**
@@ -114,7 +121,7 @@ public StarTreeValues(
         SegmentReadState readState
     ) throws IOException {
 
-        StarTreeMetadata starTreeMetadata = (StarTreeMetadata) compositeIndexMetadata;
+        starTreeMetadata = (StarTreeMetadata) compositeIndexMetadata;
 
         // build skip star node dimensions
         Set<String> skipStarNodeCreationInDims = starTreeMetadata.getSkipStarNodeCreationInDims();
@@ -262,4 +269,7 @@ public DocIdSetIterator getMetricDocIdSetIterator(String fullyQualifiedMetricNam
         return DocValues.emptySortedNumeric();
     }
 
+    public StarTreeMetadata getStarTreeMetadata() {
+        return starTreeMetadata;
+    }
 }
diff --git a/...rc/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java b/...rc/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java
@@ -37,7 +37,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 public class StarTreeTestUtils {
@@ -208,11 +208,7 @@ public static void validateFileFormats(
                         if (child.getStarTreeNodeType() != StarTreeNodeType.NULL.getValue()) {
                             assertNotNull(starTreeNode.getChildForDimensionValue(child.getDimensionValue()));
                         } else {
-                            StarTreeNode finalStarTreeNode = starTreeNode;
-                            assertThrows(
-                                AssertionError.class,
-                                () -> finalStarTreeNode.getChildForDimensionValue(child.getDimensionValue())
-                            );
+                            assertNull(starTreeNode.getChildForDimensionValue(child.getDimensionValue()));
                         }
                         assertStarTreeNode(child, resultChildNode);
                         assertNotEquals(child.getStarTreeNodeType(), StarTreeNodeType.STAR.getValue());

diff --git a/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java
@@ -2527,7 +2527,8 @@ private StarTreeValues getStarTreeValues(
             null,
             dimDocIdSetIterators,
             metricDocIdSetIterators,
-            Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, number)
+            Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, number),
+            null
         );
         return starTreeValues;
     }
@@ -3678,7 +3679,14 @@ private StarTreeValues getStarTreeValues(
         );
         // metricDocIdSetIterators.put("field2", () -> m1sndv);
         // metricDocIdSetIterators.put("_doc_count", () -> m2sndv);
-        StarTreeValues starTreeValues = new StarTreeValues(sf, null, dimDocIdSetIterators, metricDocIdSetIterators, getAttributes(500));
+        StarTreeValues starTreeValues = new StarTreeValues(
+            sf,
+            null,
+            dimDocIdSetIterators,
+            metricDocIdSetIterators,
+            getAttributes(500),
+            null
+        );
         return starTreeValues;
     }
 
@@ -4095,7 +4103,8 @@ public void testMergeFlow() throws IOException {
             null,
             dimDocIdSetIterators,
             metricDocIdSetIterators,
-            getAttributes(1000)
+            getAttributes(1000),
+            null
         );
 
         SortedNumericDocValues f2d1sndv = getSortedNumericMock(dimList1, docsWithField1);
@@ -4121,7 +4130,8 @@ public void testMergeFlow() throws IOException {
             null,
             f2dimDocIdSetIterators,
             f2metricDocIdSetIterators,
-            getAttributes(1000)
+            getAttributes(1000),
+            null
         );
 
         this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec(

diff --git a/...index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java b/...index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java
@@ -191,7 +191,7 @@ public void testGetChildForNullNode() throws IOException {
 
     public void testGetChildForInvalidDimensionValue() throws IOException {
         long invalidDimensionValue = Long.MAX_VALUE;
-        assertThrows(AssertionError.class, () -> starTreeNode.getChildForDimensionValue(invalidDimensionValue));
+        assertNull(starTreeNode.getChildForDimensionValue(invalidDimensionValue));
     }
 
     public void testOnlyRootNodePresent() throws IOException {