interfaces fixes for star tree search (#15603)

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

Hide file tree

Showing 6 changed files with 46 additions and 19 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
@@ -9,7 +9,6 @@
 package org.opensearch.index.compositeindex.datacube.startree.index;
 
 import org.apache.lucene.codecs.DocValuesProducer;
-import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SortedNumericDocValues;
@@ -74,6 +73,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 +93,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 +120,7 @@ public StarTreeValues(
         SegmentReadState readState
     ) throws IOException {
 
-        StarTreeMetadata starTreeMetadata = (StarTreeMetadata) compositeIndexMetadata;
+        starTreeMetadata = (StarTreeMetadata) compositeIndexMetadata;
 
         // build skip star node dimensions
         Set<String> skipStarNodeCreationInDims = starTreeMetadata.getSkipStarNodeCreationInDims();
@@ -244,7 +250,7 @@ public DocIdSetIterator getDimensionDocIdSetIterator(String dimension) {
             return dimensionDocValuesIteratorMap.get(dimension).get();
         }
 
-        return DocValues.emptySortedNumeric();
+        throw new IllegalArgumentException("dimension [" + dimension + "] does not exist in the segment.");
     }
 
     /**
@@ -259,7 +265,10 @@ public DocIdSetIterator getMetricDocIdSetIterator(String fullyQualifiedMetricNam
             return metricDocValuesIteratorMap.get(fullyQualifiedMetricName).get();
         }
 
-        return DocValues.emptySortedNumeric();
+        throw new IllegalArgumentException("metric [" + fullyQualifiedMetricName + "] does not exist in the segment.");
     }
 
+    public int getStarTreeDocumentCount() {
+        return starTreeMetadata.getStarTreeDocCount();
+    }
 }
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
@@ -1894,6 +1894,7 @@ private void validateStarTreeFileFormats(
         IndexInput metaIn = readState.directory.openInput(metaFileName, IOContext.DEFAULT);
 
         StarTreeValues starTreeValues = new StarTreeValues(expectedStarTreeMetadata, dataIn, compositeDocValuesProducer, readState);
+        assertEquals(expectedStarTreeMetadata.getStarTreeDocCount(), starTreeValues.getStarTreeDocumentCount());
         List<StarTreeNumericType> starTreeNumericTypes = new ArrayList<>();
         builder.metricAggregatorInfos.forEach(
             metricAggregatorInfo -> starTreeNumericTypes.add(metricAggregatorInfo.getValueAggregators().getAggregatedValueType())
@@ -2527,7 +2528,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 +3680,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;
     }
 
@@ -4088,14 +4097,20 @@ public void testMergeFlow() throws IOException {
             () -> d4sndv
         );
 
-        Map<String, Supplier<DocIdSetIterator>> metricDocIdSetIterators = Map.of("field2", () -> m1sndv, "_doc_count", () -> m2sndv);
+        Map<String, Supplier<DocIdSetIterator>> metricDocIdSetIterators = Map.of(
+            "sf_field2_sum_metric",
+            () -> m1sndv,
+            "sf__doc_count_doc_count_metric",
+            () -> m2sndv
+        );
 
         StarTreeValues starTreeValues = new StarTreeValues(
             compositeField,
             null,
             dimDocIdSetIterators,
             metricDocIdSetIterators,
-            getAttributes(1000)
+            getAttributes(1000),
+            null
         );
 
         SortedNumericDocValues f2d1sndv = getSortedNumericMock(dimList1, docsWithField1);
@@ -4115,13 +4130,19 @@ public void testMergeFlow() throws IOException {
             () -> f2d4sndv
         );
 
-        Map<String, Supplier<DocIdSetIterator>> f2metricDocIdSetIterators = Map.of("field2", () -> f2m1sndv, "_doc_count", () -> f2m2sndv);
+        Map<String, Supplier<DocIdSetIterator>> f2metricDocIdSetIterators = Map.of(
+            "sf_field2_sum_metric",
+            () -> f2m1sndv,
+            "sf__doc_count_doc_count_metric",
+            () -> f2m2sndv
+        );
         StarTreeValues starTreeValues2 = new StarTreeValues(
             compositeField,
             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 {