From d1d6c857881283d61685ce701d67ad6379a3a871 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Wed, 17 Jul 2024 18:53:15 +0530 Subject: [PATCH] rebase of file formats with merge Signed-off-by: Sarthak Aggarwal --- .../composite/Composite99DocValuesWriter.java | 1 + .../aggregators/MaxValueAggregator.java | 4 +- .../aggregators/MinValueAggregator.java | 4 +- .../startree/builder/BaseStarTreeBuilder.java | 43 +++++-------------- .../builder/OnHeapStarTreeBuilder.java | 39 +---------------- .../utils/SequentialDocValuesIterator.java | 43 +++++++++++++++++++ .../startree/utils/StarTreeBuilderUtils.java | 2 - .../utils/StarTreeDataSerializer.java | 12 +++--- .../StarTreeDocValuesFormatTests.java | 12 ++++++ .../aggregators/MaxValueAggregatorTests.java | 2 +- .../aggregators/MinValueAggregatorTests.java | 2 +- .../builder/OnHeapStarTreeBuilderTests.java | 38 ++++++++-------- .../builder/StarTreesBuilderTests.java | 12 ++++-- 13 files changed, 108 insertions(+), 106 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java b/server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java index c3f74706a346d..c3a924b34bc98 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java @@ -26,6 +26,7 @@ import org.opensearch.index.compositeindex.datacube.startree.builder.StarTreesBuilder; import org.opensearch.index.mapper.CompositeMappedFieldType; import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.StarTreeMapper; import java.io.IOException; import java.util.ArrayList; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java index e80d1c9097f9d..2c3571ae3893c 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java @@ -65,9 +65,9 @@ public Long toLongValue(Double value) { } @Override - public Double toStarTreeNumericTypeValue(Long value, StarTreeNumericType type) { + public Double toStarTreeNumericTypeValue(Long value) { try { - return type.getDoubleValue(value); + return VALUE_AGGREGATOR_TYPE.getDoubleValue(value); } catch (Exception e) { throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e); } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java index 26acf44e778a7..8c8529ff130cc 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java @@ -65,9 +65,9 @@ public Long toLongValue(Double value) { } @Override - public Double toStarTreeNumericTypeValue(Long value, StarTreeNumericType type) { + public Double toStarTreeNumericTypeValue(Long value) { try { - return type.getDoubleValue(value); + return VALUE_AGGREGATOR_TYPE.getDoubleValue(value); } catch (Exception e) { throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e); } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java index ffaa90cbaafda..da538408cc03b 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java @@ -37,6 +37,7 @@ import org.opensearch.index.compositeindex.datacube.startree.aggregators.ValueAggregator; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeBuilderUtils; import org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode; import org.opensearch.index.fielddata.IndexNumericFieldData; import org.opensearch.index.mapper.Mapper; @@ -176,12 +177,17 @@ public List getMetricReaders(SegmentWriteState stat List metricReaders = new ArrayList<>(); for (Metric metric : this.starTreeField.getMetrics()) { - for (MetricStat metricType : metric.getMetrics()) { + for (MetricStat metricStat : metric.getMetrics()) { SequentialDocValuesIterator metricReader; FieldInfo metricFieldInfo = state.fieldInfos.fieldInfo(metric.getField()); + if (metricStat != MetricStat.COUNT) { metricReader = new SequentialDocValuesIterator( - fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo) - ); + fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo) + ); + } else { + metricReader = new SequentialDocValuesIterator(); + } + metricReaders.add(metricReader); } } @@ -215,7 +221,6 @@ public void build( ); } Iterator starTreeDocumentIterator = sortAndAggregateSegmentDocuments( - totalSegmentDocs, dimensionReaders, metricReaders ); @@ -463,7 +468,6 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) { * Sorts and aggregates all the documents in the segment as per the configuration, and returns a star-tree document iterator for all the * aggregated star-tree documents. * - * @param numDocs number of documents in the given segment * @param dimensionReaders List of docValues readers to read dimensions from the segment * @param metricReaders List of docValues readers to read metrics from the segment * @return Iterator for the aggregated star-tree document @@ -694,12 +698,12 @@ public void build(Map fieldProducerMap) throws IOExce logger.debug("No documents found in the segment"); return; } - List metricReaders = getMetricReaders(state, fieldProducerMap); + List metricReaders = getMetricReaders(writeState, fieldProducerMap); List dimensionsSplitOrder = starTreeField.getDimensionsOrder(); SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[dimensionsSplitOrder.size()]; for (int i = 0; i < numDimensions; i++) { String dimension = dimensionsSplitOrder.get(i).getField(); - FieldInfo dimensionFieldInfo = state.fieldInfos.fieldInfo(dimension); + FieldInfo dimensionFieldInfo = writeState.fieldInfos.fieldInfo(dimension); dimensionReaders[i] = new SequentialDocValuesIterator( fieldProducerMap.get(dimensionFieldInfo.name).getSortedNumeric(dimensionFieldInfo) ); @@ -710,31 +714,6 @@ public void build(Map fieldProducerMap) throws IOExce logger.debug("Finished Building star-tree in ms : {}", (System.currentTimeMillis() - startTime)); } - /** - * Generates the configuration required to perform aggregation for all the metrics on a field - * - * @return list of MetricAggregatorInfo - */ - public List getMetricReaders(SegmentWriteState state, Map fieldProducerMap) - throws IOException { - List metricReaders = new ArrayList<>(); - for (Metric metric : this.starTreeField.getMetrics()) { - for (MetricStat metricType : metric.getMetrics()) { - FieldInfo metricFieldInfo = state.fieldInfos.fieldInfo(metric.getField()); - // TODO - // if (metricType != MetricStat.COUNT) { - // Need not initialize the metric reader for COUNT metric type - SequentialDocValuesIterator metricReader = new SequentialDocValuesIterator( - fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo) - ); - // } - - metricReaders.add(metricReader); - } - } - return metricReaders; - } - /** * Builds the star tree using Star-Tree Document * diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java index a2be7a8b1c0c3..35a03227114c3 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java @@ -92,7 +92,7 @@ public Iterator sortAndAggregateSegmentDocuments( for (int currentDocId = 0; currentDocId < totalSegmentDocs; currentDocId++) { starTreeDocuments[currentDocId] = getSegmentStarTreeDocument(currentDocId, dimensionReaders, metricReaders); } - return sortAndAggregateStarTreeDocuments(starTreeDocuments); + return sortAndAggregateStarTreeDocuments(starTreeDocuments, false); } @Override @@ -179,48 +179,11 @@ StarTreeDocument[] mergeStarTreeValues(List starTreeValuesSubs) return starTreeDocuments.toArray(starTreeDocumentsArr); } - @Override - public StarTreeDocument getStarTreeDocument(int docId) throws IOException { - return starTreeDocuments.get(docId); - } - @Override public StarTreeDocument getStarTreeDocumentForCreatingDocValues(int docId) throws IOException { return starTreeDocuments.get(docId); } - @Override - public List getStarTreeDocuments() { - return starTreeDocuments; - } - - @Override - public Long getDimensionValue(int docId, int dimensionId) throws IOException { - return starTreeDocuments.get(docId).dimensions[dimensionId]; - } - - /** - * Sorts and aggregates all the documents of the segment based on dimension and metrics configuration - * - * @param numDocs number of documents in the given segment - * @param dimensionReaders List of docValues readers to read dimensions from the segment - * @param metricReaders List of docValues readers to read metrics from the segment - * @return Iterator of star-tree documents - * - */ - @Override - public Iterator sortAndAggregateSegmentDocuments( - int numDocs, - SequentialDocValuesIterator[] dimensionReaders, - List metricReaders - ) throws IOException { - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[numDocs]; - for (int currentDocId = 0; currentDocId < numDocs; currentDocId++) { - starTreeDocuments[currentDocId] = getSegmentStarTreeDocument(currentDocId, dimensionReaders, metricReaders); - } - return sortAndAggregateStarTreeDocuments(starTreeDocuments); - } - /** * Sort, aggregates and merges the star-tree documents * diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java index 5e030233c5e8b..b3360571c9dfc 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java @@ -47,6 +47,49 @@ public SequentialDocValuesIterator(DocIdSetIterator docIdSetIterator) { this.docIdSetIterator = docIdSetIterator; } + /** + * Constructs a new SequentialDocValuesIterator instance with the given SortedNumericDocValues. + * + */ + public SequentialDocValuesIterator() { + this.docIdSetIterator = new SortedNumericDocValues() { + @Override + public long nextValue() throws IOException { + return 0; + } + + @Override + public int docValueCount() { + return 0; + } + + @Override + public boolean advanceExact(int i) throws IOException { + return false; + } + + @Override + public int docID() { + return 0; + } + + @Override + public int nextDoc() throws IOException { + return 0; + } + + @Override + public int advance(int i) throws IOException { + return 0; + } + + @Override + public long cost() { + return 0; + } + }; + } + /** * Returns the value associated with the latest document. * diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeBuilderUtils.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeBuilderUtils.java index cb0fc100db546..d22c75be74ea4 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeBuilderUtils.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeBuilderUtils.java @@ -9,14 +9,12 @@ import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.store.IndexOutput; -import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo; import org.opensearch.index.mapper.CompositeMappedFieldType; import java.io.IOException; import java.util.List; -import java.util.Map; /** * Util class for building star tree diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDataSerializer.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDataSerializer.java index 54059458a6d0d..898234cdacc86 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDataSerializer.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDataSerializer.java @@ -42,7 +42,7 @@ public class StarTreeDataSerializer { * @return the total size in bytes of the serialized star-tree data * @throws IOException if an I/O error occurs while writing the star-tree data */ - public static long serializeStarTree(IndexOutput indexOutput, StarTreeBuilderUtils.TreeNode rootNode, int numNodes) throws IOException { + public static long serializeStarTree(IndexOutput indexOutput, TreeNode rootNode, int numNodes) throws IOException { int headerSizeInBytes = computeStarTreeDataHeaderByteSize(); long totalSizeInBytes = headerSizeInBytes + (long) numNodes * SERIALIZABLE_DATA_SIZE_IN_BYTES; @@ -87,20 +87,20 @@ private static void writeStarTreeHeader(IndexOutput output, int numNodes) throws * @param rootNode the root node of the star-tree * @throws IOException if an I/O error occurs while writing the nodes */ - private static void writeStarTreeNodes(IndexOutput output, StarTreeBuilderUtils.TreeNode rootNode) throws IOException { - Queue queue = new LinkedList<>(); + private static void writeStarTreeNodes(IndexOutput output, TreeNode rootNode) throws IOException { + Queue queue = new LinkedList<>(); queue.add(rootNode); int currentNodeId = 0; while (!queue.isEmpty()) { - StarTreeBuilderUtils.TreeNode node = queue.remove(); + TreeNode node = queue.remove(); if (node.children == null) { writeStarTreeNode(output, node, ALL, ALL); } else { // Sort all children nodes based on dimension value - List sortedChildren = new ArrayList<>(node.children.values()); + List sortedChildren = new ArrayList<>(node.children.values()); sortedChildren.sort(Comparator.comparingLong(o -> o.dimensionValue)); int firstChildId = currentNodeId + queue.size() + 1; @@ -123,7 +123,7 @@ private static void writeStarTreeNodes(IndexOutput output, StarTreeBuilderUtils. * @param lastChildId the ID of the last child node * @throws IOException if an I/O error occurs while writing the node */ - private static void writeStarTreeNode(IndexOutput output, StarTreeBuilderUtils.TreeNode node, int firstChildId, int lastChildId) + private static void writeStarTreeNode(IndexOutput output, TreeNode node, int firstChildId, int lastChildId) throws IOException { output.writeInt(node.dimensionId); output.writeLong(node.dimensionValue); diff --git a/server/src/test/java/org/opensearch/index/codec/composite/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite/datacube/startree/StarTreeDocValuesFormatTests.java index 9a7ae7525e5db..4d07d6afe0323 100644 --- a/server/src/test/java/org/opensearch/index/codec/composite/datacube/startree/StarTreeDocValuesFormatTests.java +++ b/server/src/test/java/org/opensearch/index/codec/composite/datacube/startree/StarTreeDocValuesFormatTests.java @@ -23,6 +23,7 @@ import org.opensearch.cluster.ClusterModule; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.CheckedConsumer; +import org.opensearch.common.Rounding; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.FeatureFlags; import org.opensearch.common.xcontent.XContentFactory; @@ -30,14 +31,25 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.MapperTestUtils; import org.opensearch.index.codec.composite.Composite99Codec; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.StarTreeMapper; import org.opensearch.indices.IndicesModule; import org.junit.After; import org.junit.AfterClass; import org.junit.BeforeClass; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.DateDimension; +import org.opensearch.index.compositeindex.datacube.NumericDimension; + import java.io.IOException; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX; import static org.opensearch.test.OpenSearchTestCase.randomFrom; diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java index 194d667cffbd7..1ff600643f198 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java @@ -53,6 +53,6 @@ public void testToLongValue() { } public void testToStarTreeNumericTypeValue() { - assertEquals(NumericUtils.sortableLongToDouble(3L), aggregator.toStarTreeNumericTypeValue(3L, StarTreeNumericType.DOUBLE), 0.0); + assertEquals(NumericUtils.sortableLongToDouble(3L), aggregator.toStarTreeNumericTypeValue(3L), 0.0); } } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java index 8f8ce259bc62c..0257fbd3ea71c 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java @@ -53,6 +53,6 @@ public void testToLongValue() { } public void testToStarTreeNumericTypeValue() { - assertEquals(NumericUtils.sortableLongToDouble(3L), aggregator.toStarTreeNumericTypeValue(3L, StarTreeNumericType.DOUBLE), 0.0); + assertEquals(NumericUtils.sortableLongToDouble(3L), aggregator.toStarTreeNumericTypeValue(3L), 0.0); } } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java index 0d5dc3d9ae77d..d79d2ced98492 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java @@ -214,7 +214,7 @@ public void test_sortAndAggregateStarTreeDocuments() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); int numOfAggregatedDocuments = 0; while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); @@ -258,7 +258,7 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetric() throws IOExcepti segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); for (int dim = 0; dim < 4; dim++) { @@ -318,7 +318,7 @@ public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); int numOfAggregatedDocuments = 0; while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); @@ -371,7 +371,7 @@ public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); int numOfAggregatedDocuments = 0; while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); @@ -482,7 +482,7 @@ public void test_build_halfFloatMetrics() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); List resultStarTreeDocuments = builder.getStarTreeDocuments(); assertEquals(7, resultStarTreeDocuments.size()); @@ -547,7 +547,7 @@ public void test_build_floatMetrics() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); List resultStarTreeDocuments = builder.getStarTreeDocuments(); @@ -601,7 +601,7 @@ public void test_build_longMetrics() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); List resultStarTreeDocuments = builder.getStarTreeDocuments(); @@ -613,13 +613,13 @@ public void test_build_longMetrics() throws IOException { private static Iterator getExpectedStarTreeDocumentIterator() { List expectedStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L, 6.0, 24.0 }) ); return expectedStarTreeDocuments.iterator(); } @@ -644,7 +644,7 @@ public void test_build() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, metrics); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); List resultStarTreeDocuments = builder.getStarTreeDocuments(); @@ -746,7 +746,7 @@ public void test_build_starTreeDataset() throws IOException { null ); when(documentMapper.mappers()).thenReturn(fieldMappers); - builder = new OnHeapStarTreeBuilder(compositeField, writeState, mapperService); + builder = new OnHeapStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); int noOfStarTreeDocuments = 7; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; @@ -764,7 +764,7 @@ public void test_build_starTreeDataset() throws IOException { segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1 }); } - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateStarTreeDocuments(segmentStarTreeDocuments, false); builder.build(segmentStarTreeDocumentIterator); List resultStarTreeDocuments = builder.getStarTreeDocuments(); @@ -845,7 +845,7 @@ public void testFlushFlow() throws IOException { SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); SortedNumericDocValues m2sndv = getSortedNumericMock(metricsList, metricsWithField); - OnHeapStarTreeBuilder builder = new OnHeapStarTreeBuilder(sf, writeState, mapperService); + OnHeapStarTreeBuilder builder = new OnHeapStarTreeBuilder(metaOut, dataOut, sf, writeState, mapperService); SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( dimDvs, @@ -970,7 +970,7 @@ public void testMergeFlowWithCount() throws IOException { Map f2dimDocIdSetIterators = Map.of("field1", f2d1sndv, "field3", f2d2sndv); Map f2metricDocIdSetIterators = Map.of("field2", f2m1sndv); StarTreeValues starTreeValues2 = new StarTreeValues(sf, null, f2dimDocIdSetIterators, f2metricDocIdSetIterators); - OnHeapStarTreeBuilder builder = new OnHeapStarTreeBuilder(sf, writeState, mapperService); + OnHeapStarTreeBuilder builder = new OnHeapStarTreeBuilder(metaOut, dataOut, sf, writeState, mapperService); Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java index 564ab110fa7a5..455cd84bb44f8 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube.startree.builder; +import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.apache.lucene.index.FieldInfo; @@ -15,6 +16,7 @@ import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.Version; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; @@ -45,9 +47,13 @@ public class StarTreesBuilderTests extends OpenSearchTestCase { private StarTreeField starTreeField; private Map fieldProducerMap; private Directory directory; + private IndexOutput dataOut; + private IndexOutput metaOut; public void setUp() throws Exception { super.setUp(); + metaOut = mock(IndexOutput.class); + dataOut = mock(IndexOutput.class); mapperService = mock(MapperService.class); directory = newFSDirectory(createTempDir()); SegmentInfo segmentInfo = new SegmentInfo( @@ -89,7 +95,7 @@ public void test_buildWithNoStarTreeFields() throws IOException { when(mapperService.getCompositeFieldTypes()).thenReturn(new HashSet<>()); StarTreesBuilder starTreesBuilder = new StarTreesBuilder(segmentWriteState, mapperService); - starTreesBuilder.build(fieldProducerMap); + starTreesBuilder.build(metaOut, dataOut, fieldProducerMap, mock(DocValuesConsumer.class)); verifyNoInteractions(docValuesProducer); } @@ -97,7 +103,7 @@ public void test_buildWithNoStarTreeFields() throws IOException { public void test_getStarTreeBuilder() throws IOException { when(mapperService.getCompositeFieldTypes()).thenReturn(Set.of(starTreeFieldType)); StarTreesBuilder starTreesBuilder = new StarTreesBuilder(segmentWriteState, mapperService); - StarTreeBuilder starTreeBuilder = starTreesBuilder.getSingleTreeBuilder(starTreeField, segmentWriteState, mapperService); + StarTreeBuilder starTreeBuilder = starTreesBuilder.getSingleTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService); assertTrue(starTreeBuilder instanceof OnHeapStarTreeBuilder); } @@ -106,7 +112,7 @@ public void test_getStarTreeBuilder_illegalArgument() { StarTreeFieldConfiguration starTreeFieldConfiguration = new StarTreeFieldConfiguration(1, new HashSet<>(), StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP); StarTreeField starTreeField = new StarTreeField("star_tree", new ArrayList<>(), new ArrayList<>(), starTreeFieldConfiguration); StarTreesBuilder starTreesBuilder = new StarTreesBuilder(segmentWriteState, mapperService); - assertThrows(IllegalArgumentException.class, () -> starTreesBuilder.getSingleTreeBuilder(starTreeField, segmentWriteState, mapperService)); + assertThrows(IllegalArgumentException.class, () -> starTreesBuilder.getSingleTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService)); } public void test_closeWithNoStarTreeFields() throws IOException {