From 08722e1bd7b2e486aaf95176a37f5562098dbf44 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Tue, 3 Sep 2024 13:17:16 +0530 Subject: [PATCH] interfaces fixes for star tree search Signed-off-by: Sarthak Aggarwal --- .../fileformats/meta/StarTreeMetadata.java | 2 ++ .../node/FixedLengthStarTreeNode.java | 1 - .../startree/index/StarTreeValues.java | 19 ++++++++--- .../datacube/startree/StarTreeTestUtils.java | 8 ++--- .../builder/AbstractStarTreeBuilderTests.java | 33 +++++++++++++++---- .../node/FixedLengthStarTreeNodeTests.java | 2 +- 6 files changed, 46 insertions(+), 19 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java index 2e20c9b9e0737..e7bb32282ece7 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadata.java +++ b/server/src/main/java/org/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/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java index 72c08d96d26ff..0ae704afa53b9 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java +++ b/server/src/main/java/org/opensearch/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/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java index 8181a69f49bd4..7d5f5ba02b9f8 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java +++ b/server/src/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 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> dimensionDocValuesIteratorMap, Map> metricDocValuesIteratorMap, - Map attributes + Map 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 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/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java index 53c9e05ccb490..7b5672529dc8d 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java +++ b/server/src/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/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java index 79d1b0e4b9785..6e23e5eacaecd 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java +++ b/server/src/test/java/org/opensearch/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 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> metricDocIdSetIterators = Map.of("field2", () -> m1sndv, "_doc_count", () -> m2sndv); + Map> 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> f2metricDocIdSetIterators = Map.of("field2", () -> f2m1sndv, "_doc_count", () -> f2m2sndv); + Map> 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/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java index 9657f867422a0..08815d5ef55f5 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java +++ b/server/src/test/java/org/opensearch/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 {