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 ddcf02cc6291a..d4bb1cc3dc1f7 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 @@ -26,8 +26,9 @@ import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo; import org.opensearch.index.compositeindex.datacube.startree.aggregators.ValueAggregator; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; -import org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode; import org.opensearch.index.fielddata.IndexNumericFieldData; import org.opensearch.index.mapper.DocCountFieldMapper; import org.opensearch.index.mapper.Mapper; @@ -46,7 +47,7 @@ import java.util.Objects; import java.util.Set; -import static org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode.ALL; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL; /** * Builder for star tree. Defines the algorithm to construct star-tree @@ -72,7 +73,7 @@ public abstract class BaseStarTreeBuilder implements StarTreeBuilder { protected int numStarTreeNodes; protected final int maxLeafDocuments; - protected final TreeNode rootNode = getNewNode(); + protected final InMemoryTreeNode rootNode = getNewNode(); protected final StarTreeField starTreeField; private final SegmentWriteState state; @@ -578,9 +579,9 @@ private void appendToStarTree(StarTreeDocument starTreeDocument) throws IOExcept * * @return return new star-tree node */ - private TreeNode getNewNode() { + private InMemoryTreeNode getNewNode() { numStarTreeNodes++; - return new TreeNode(); + return new InMemoryTreeNode(); } /** @@ -591,7 +592,7 @@ private TreeNode getNewNode() { * @param endDocId end document id * @throws IOException throws an exception if we are unable to construct the tree */ - private void constructStarTree(TreeNode node, int startDocId, int endDocId) throws IOException { + private void constructStarTree(InMemoryTreeNode node, int startDocId, int endDocId) throws IOException { int childDimensionId = node.dimensionId + 1; if (childDimensionId == numDimensions) { @@ -600,7 +601,7 @@ private void constructStarTree(TreeNode node, int startDocId, int endDocId) thro // Construct all non-star children nodes node.childDimensionId = childDimensionId; - Map children = constructNonStarNodes(startDocId, endDocId, childDimensionId); + Map children = constructNonStarNodes(startDocId, endDocId, childDimensionId); node.children = children; // Construct star-node if required @@ -609,7 +610,7 @@ private void constructStarTree(TreeNode node, int startDocId, int endDocId) thro } // Further split on child nodes if required - for (TreeNode child : children.values()) { + for (InMemoryTreeNode child : children.values()) { if (child.endDocId - child.startDocId > maxLeafDocuments) { constructStarTree(child, child.startDocId, child.endDocId); } @@ -625,14 +626,14 @@ private void constructStarTree(TreeNode node, int startDocId, int endDocId) thro * @return root node with non-star nodes constructed * @throws IOException throws an exception if we are unable to construct non-star nodes */ - private Map constructNonStarNodes(int startDocId, int endDocId, int dimensionId) throws IOException { - Map nodes = new HashMap<>(); + private Map constructNonStarNodes(int startDocId, int endDocId, int dimensionId) throws IOException { + Map nodes = new HashMap<>(); int nodeStartDocId = startDocId; Long nodeDimensionValue = getDimensionValue(startDocId, dimensionId); for (int i = startDocId + 1; i < endDocId; i++) { Long dimensionValue = getDimensionValue(i, dimensionId); if (Objects.equals(dimensionValue, nodeDimensionValue) == false) { - TreeNode child = getNewNode(); + InMemoryTreeNode child = getNewNode(); child.dimensionId = dimensionId; child.dimensionValue = nodeDimensionValue != null ? nodeDimensionValue : ALL; child.startDocId = nodeStartDocId; @@ -643,7 +644,7 @@ private Map constructNonStarNodes(int startDocId, int endDocId, nodeDimensionValue = dimensionValue; } } - TreeNode lastNode = getNewNode(); + InMemoryTreeNode lastNode = getNewNode(); lastNode.dimensionId = dimensionId; lastNode.dimensionValue = nodeDimensionValue != null ? nodeDimensionValue : ALL; lastNode.startDocId = nodeStartDocId; @@ -661,11 +662,11 @@ private Map constructNonStarNodes(int startDocId, int endDocId, * @return root node with star nodes constructed * @throws IOException throws an exception if we are unable to construct non-star nodes */ - private TreeNode constructStarNode(int startDocId, int endDocId, int dimensionId) throws IOException { - TreeNode starNode = getNewNode(); + private InMemoryTreeNode constructStarNode(int startDocId, int endDocId, int dimensionId) throws IOException { + InMemoryTreeNode starNode = getNewNode(); starNode.dimensionId = dimensionId; starNode.dimensionValue = ALL; - starNode.isStarNode = true; + starNode.nodeType = StarTreeNodeType.STAR.getValue(); starNode.startDocId = numStarTreeDocs; Iterator starTreeDocumentIterator = generateStarTreeDocumentsForStarNode(startDocId, endDocId, dimensionId); while (starTreeDocumentIterator.hasNext()) { @@ -682,7 +683,7 @@ private TreeNode constructStarNode(int startDocId, int endDocId, int dimensionId * @return aggregated star-tree documents * @throws IOException throws an exception upon failing to create new aggregated docs based on star tree */ - private StarTreeDocument createAggregatedDocs(TreeNode node) throws IOException { + private StarTreeDocument createAggregatedDocs(InMemoryTreeNode node) throws IOException { StarTreeDocument aggregatedStarTreeDocument = null; if (node.children == null) { @@ -709,8 +710,8 @@ private StarTreeDocument createAggregatedDocs(TreeNode node) throws IOException // For non-leaf node if (node.children.containsKey((long) ALL)) { // If it has star child, use the star child aggregated document directly - for (TreeNode child : node.children.values()) { - if (child.isStarNode) { + for (InMemoryTreeNode child : node.children.values()) { + if (child.nodeType == StarTreeNodeType.STAR.getValue()) { aggregatedStarTreeDocument = createAggregatedDocs(child); node.aggregatedDocId = child.aggregatedDocId; } else { @@ -720,12 +721,12 @@ private StarTreeDocument createAggregatedDocs(TreeNode node) throws IOException } else { // If no star child exists, aggregate all aggregated documents from non-star children if (node.children.values().size() == 1) { - for (TreeNode child : node.children.values()) { + for (InMemoryTreeNode child : node.children.values()) { aggregatedStarTreeDocument = reduceStarTreeDocuments(aggregatedStarTreeDocument, createAggregatedDocs(child)); node.aggregatedDocId = child.aggregatedDocId; } } else { - for (TreeNode child : node.children.values()) { + for (InMemoryTreeNode child : node.children.values()) { aggregatedStarTreeDocument = reduceStarTreeDocuments(aggregatedStarTreeDocument, createAggregatedDocs(child)); } if (null == aggregatedStarTreeDocument) { @@ -760,7 +761,7 @@ public void close() throws IOException { abstract Iterator mergeStarTrees(List starTreeValues) throws IOException; - public TreeNode getRootNode() { + public InMemoryTreeNode getRootNode() { return rootNode; } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java index 6f11148dee468..32feb78a4db3d 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeDataWriter.java @@ -20,7 +20,7 @@ import java.util.List; import java.util.Queue; -import static org.opensearch.index.compositeindex.datacube.startree.node.FixedLengthStarTreeNode.SERIALIZABLE_DATA_SIZE_IN_BYTES; +import static org.opensearch.index.compositeindex.datacube.startree.fileformats.node.FixedLengthStarTreeNode.SERIALIZABLE_DATA_SIZE_IN_BYTES; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL; /** diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/MetricEntry.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/MetricEntry.java deleted file mode 100644 index 357c8a49f600c..0000000000000 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/MetricEntry.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.index.compositeindex.datacube.startree.fileformats.meta; - -import org.opensearch.index.compositeindex.datacube.MetricStat; - -import java.util.Objects; - -/** - * Holds the pair of metric name and it's associated stat - * - * @opensearch.experimental - */ -public class MetricEntry { - - private final String metricFieldName; - private final MetricStat metricStat; - - public MetricEntry(String metricFieldName, MetricStat metricStat) { - this.metricFieldName = metricFieldName; - this.metricStat = metricStat; - } - - public String getMetricFieldName() { - return metricFieldName; - } - - public MetricStat getMetricStat() { - return metricStat; - } - - @Override - public int hashCode() { - return Objects.hashCode(metricFieldName + metricStat.getTypeName()); - } - - @Override - public boolean equals(Object obj) { - if (this == obj) { - return true; - } - if (obj instanceof MetricEntry) { - MetricEntry anotherPair = (MetricEntry) obj; - return metricStat.equals(anotherPair.metricStat) && metricFieldName.equals(anotherPair.metricFieldName); - } - return false; - } - -} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaWriter.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaWriter.java index f7ff91f5928af..112e47d50fde6 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaWriter.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaWriter.java @@ -119,9 +119,11 @@ private static void writeMeta( metaOut.writeInt(numNodes); // number of dimensions + // TODO: Revisit the number of dimensions for timestamps (as we will split timestamp into min, hour, etc.) metaOut.writeVInt(starTreeField.getDimensionsOrder().size()); // dimensions + // TODO: Add sub-dimensions for timestamps (as we will split timestamp into min, hour, etc.) for (Dimension dimension : starTreeField.getDimensionsOrder()) { metaOut.writeString(dimension.getField()); } 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 e8d24cf8f6209..7519c85562a8c 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 @@ -12,6 +12,7 @@ import org.apache.lucene.index.CorruptIndexException; import org.apache.lucene.store.IndexInput; import org.opensearch.index.compositeindex.CompositeIndexMetadata; +import org.opensearch.index.compositeindex.datacube.Metric; import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.mapper.CompositeMappedFieldType; @@ -19,7 +20,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Set; /** @@ -62,9 +65,9 @@ public class StarTreeMetadata extends CompositeIndexMetadata { private final List dimensionFields; /** - * List of metric entries, containing field names and associated metric statistic. + * List of metrics, containing field names and associated metric statistics. */ - private final List metricEntries; + private final List metrics; /** * The total number of documents aggregated in this star-tree segment. @@ -119,7 +122,7 @@ public StarTreeMetadata( this.version = version; this.numberOfNodes = readNumberOfNodes(); this.dimensionFields = readStarTreeDimensions(); - this.metricEntries = readMetricEntries(); + this.metrics = readMetricEntries(); this.segmentAggregatedDocCount = readSegmentAggregatedDocCount(); this.maxLeafDocs = readMaxLeafDocs(); this.skipStarNodeCreationInDims = readSkipStarNodeCreationInDims(); @@ -141,7 +144,7 @@ public StarTreeMetadata( * @param compositeFieldType type of the composite field. Here, STAR_TREE field. * @param version The version of the star tree stored in the segments. * @param dimensionFields list of dimension fields - * @param metricEntries list of metric entries + * @param metrics list of metric entries * @param segmentAggregatedDocCount segment aggregated doc count * @param maxLeafDocs max leaf docs * @param skipStarNodeCreationInDims set of dimensions to skip star node creation @@ -156,7 +159,7 @@ public StarTreeMetadata( Integer version, Integer numberOfNodes, List dimensionFields, - List metricEntries, + List metrics, Integer segmentAggregatedDocCount, Integer maxLeafDocs, Set skipStarNodeCreationInDims, @@ -171,7 +174,7 @@ public StarTreeMetadata( this.version = version; this.numberOfNodes = numberOfNodes; this.dimensionFields = dimensionFields; - this.metricEntries = metricEntries; + this.metrics = metrics; this.segmentAggregatedDocCount = segmentAggregatedDocCount; this.maxLeafDocs = maxLeafDocs; this.skipStarNodeCreationInDims = skipStarNodeCreationInDims; @@ -203,17 +206,19 @@ private int readMetricsCount() throws IOException { return meta.readVInt(); } - private List readMetricEntries() throws IOException { + private List readMetricEntries() throws IOException { int metricCount = readMetricsCount(); - List metricEntries = new ArrayList<>(); + Map starTreeMetricMap = new LinkedHashMap<>(); for (int i = 0; i < metricCount; i++) { - String metricFieldName = meta.readString(); + String metricName = meta.readString(); int metricStatOrdinal = meta.readVInt(); - metricEntries.add(new MetricEntry(metricFieldName, MetricStat.fromMetricOrdinal(metricStatOrdinal))); + MetricStat metricStat = MetricStat.fromMetricOrdinal(metricStatOrdinal); + Metric metric = starTreeMetricMap.computeIfAbsent(metricName, field -> new Metric(field, new ArrayList<>())); + metric.getMetrics().add(metricStat); } - return metricEntries; + return new ArrayList<>(starTreeMetricMap.values()); } private int readSegmentAggregatedDocCount() throws IOException { @@ -282,8 +287,8 @@ public List getDimensionFields() { * * @return star-tree metric entries */ - public List getMetricEntries() { - return metricEntries; + public List getMetrics() { + return metrics; } /** diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java similarity index 94% rename from server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNode.java rename to server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java index df0662ccb1fb1..f1663de7087c4 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNode.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java @@ -5,16 +5,16 @@ * this file be licensed under the Apache-2.0 license or a * compatible open source license. */ -package org.opensearch.index.compositeindex.datacube.startree.node; +package org.opensearch.index.compositeindex.datacube.startree.fileformats.node; import org.apache.lucene.store.RandomAccessInput; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; import java.io.IOException; import java.io.UncheckedIOException; import java.util.Iterator; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL; - /** * Fixed Length implementation of {@link StarTreeNode}. *

@@ -187,16 +187,17 @@ public byte getStarTreeNodeType() throws IOException { } @Override - public StarTreeNode getChildForDimensionValue(Long dimensionValue, boolean isStar) throws IOException { + public StarTreeNode getChildStarNode() throws IOException { + return handleStarNode(); + } + + @Override + public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException { // there will be no children for leaf nodes if (isLeaf()) { return null; } - // Specialize star node for performance - if (isStar) { - return handleStarNode(); - } StarTreeNode resultStarTreeNode = null; if (null != dimensionValue) { resultStarTreeNode = binarySearchChild(dimensionValue); @@ -213,7 +214,7 @@ public StarTreeNode getChildForDimensionValue(Long dimensionValue, boolean isSta */ private FixedLengthStarTreeNode handleStarNode() throws IOException { FixedLengthStarTreeNode firstNode = new FixedLengthStarTreeNode(in, firstChildId); - if (firstNode.getDimensionValue() == ALL) { + if (firstNode.getStarTreeNodeType() == StarTreeNodeType.STAR.getValue()) { return firstNode; } else { return null; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/package-info.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/package-info.java new file mode 100644 index 0000000000000..84271be81f5e4 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Holds classes associated with star tree node with file formats + */ +package org.opensearch.index.compositeindex.datacube.startree.fileformats.node; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTree.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTree.java index 93d0cc93e8843..828d98368b9e6 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTree.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTree.java @@ -10,6 +10,7 @@ import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.RandomAccessInput; import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.node.FixedLengthStarTreeNode; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java index d29d9145853ad..fce3e30e9ebf6 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java @@ -101,13 +101,21 @@ public interface StarTreeNode { byte getStarTreeNodeType() throws IOException; /** - * Returns the child star-tree node for the given dimension value. + * Returns the child node for the given dimension value in the star-tree. * * @param dimensionValue the dimension value * @return the child node for the given dimension value or null if child is not present * @throws IOException if an I/O error occurs while retrieving the child node */ - StarTreeNode getChildForDimensionValue(Long dimensionValue, boolean isStar) throws IOException; + StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException; + + /** + * Returns the child star node for a node in the star-tree. + * + * @return the child node for the star node if star child node is not present + * @throws IOException if an I/O error occurs while retrieving the child node + */ + StarTreeNode getChildStarNode() throws IOException; /** * Returns an iterator over the children of the current star-tree node. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/TreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/TreeNode.java deleted file mode 100644 index a5d59a2602633..0000000000000 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/TreeNode.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ -package org.opensearch.index.compositeindex.datacube.startree.utils; - -import org.opensearch.common.annotation.ExperimentalApi; - -import java.util.Map; - -/** - * /** - * Represents a node in a tree data structure, specifically designed for a star-tree implementation. - * A star-tree node will represent both star and non-star nodes. - * - * @opensearch.experimental - */ -@ExperimentalApi -public class TreeNode { - - public static final int ALL = -1; - - /** - * The dimension id for the dimension (field) associated with this star-tree node. - */ - public int dimensionId = ALL; - - /** - * The starting document id (inclusive) associated with this star-tree node. - */ - public int startDocId = ALL; - - /** - * The ending document id (exclusive) associated with this star-tree node. - */ - public int endDocId = ALL; - - /** - * The aggregated document id associated with this star-tree node. - */ - public int aggregatedDocId = ALL; - - /** - * The child dimension identifier associated with this star-tree node. - */ - public int childDimensionId = ALL; - - /** - * The value of the dimension associated with this star-tree node. - */ - public long dimensionValue = ALL; - - /** - * A flag indicating whether this node is a star node (a node that represents an aggregation of all dimensions). - */ - public boolean isStarNode = false; - - /** - * A map containing the child nodes of this star-tree node, keyed by their dimension id. - */ - public Map children; - - public long getDimensionValue() { - return dimensionValue; - } -} 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 e77f184ac0243..33088e8ccbcb3 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 @@ -36,8 +36,10 @@ import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; -import org.opensearch.index.compositeindex.datacube.startree.utils.TreeNode; +import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; import org.opensearch.index.mapper.ContentPath; import org.opensearch.index.mapper.DocumentMapper; import org.opensearch.index.mapper.Mapper; @@ -1314,7 +1316,7 @@ public void test_build_starTreeDataset() throws IOException { Iterator expectedStarTreeDocumentIterator = expectedStarTreeDocuments(); Iterator resultStarTreeDocumentIterator = resultStarTreeDocuments.iterator(); Map> dimValueToDocIdMap = new HashMap<>(); - builder.rootNode.isStarNode = true; + builder.rootNode.nodeType = StarTreeNodeType.STAR.getValue(); traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); Map> expectedDimToValueMap = getExpectedDimToValueMap(); @@ -2784,13 +2786,13 @@ private static StarTreeField getStarTreeFieldWithDocCount(int maxLeafDocs, boole return sf; } - private void traverseStarTree(TreeNode root, Map> dimValueToDocIdMap, boolean traverStarNodes) { - TreeNode starTree = root; + private void traverseStarTree(InMemoryTreeNode root, Map> dimValueToDocIdMap, boolean traverStarNodes) { + InMemoryTreeNode starTree = root; // Use BFS to traverse the star tree - Queue queue = new ArrayDeque<>(); + Queue queue = new ArrayDeque<>(); queue.add(starTree); int currentDimensionId = -1; - TreeNode starTreeNode; + InMemoryTreeNode starTreeNode; List docIds = new ArrayList<>(); while ((starTreeNode = queue.poll()) != null) { int dimensionId = starTreeNode.dimensionId; @@ -2801,17 +2803,17 @@ private void traverseStarTree(TreeNode root, Map> di // store aggregated document of the node int docId = starTreeNode.aggregatedDocId; Map map = dimValueToDocIdMap.getOrDefault(dimensionId, new HashMap<>()); - if (starTreeNode.isStarNode) { + if (starTreeNode.nodeType == StarTreeNodeType.STAR.getValue()) { map.put(Long.MAX_VALUE, docId); } else { map.put(starTreeNode.dimensionValue, docId); } dimValueToDocIdMap.put(dimensionId, map); - if (starTreeNode.children != null && (!traverStarNodes || starTreeNode.isStarNode)) { - Iterator childrenIterator = starTreeNode.children.values().iterator(); + if (starTreeNode.children != null && (!traverStarNodes || starTreeNode.nodeType == StarTreeNodeType.STAR.getValue())) { + Iterator childrenIterator = starTreeNode.children.values().iterator(); while (childrenIterator.hasNext()) { - TreeNode childNode = childrenIterator.next(); + InMemoryTreeNode childNode = childrenIterator.next(); queue.add(childNode); } } @@ -2933,43 +2935,48 @@ public void testMergeFlow() throws IOException { validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); } - private void validateStarTree(TreeNode root, int totalDimensions, int maxLeafDocuments, List starTreeDocuments) { + private void validateStarTree( + InMemoryTreeNode root, + int totalDimensions, + int maxLeafDocuments, + List starTreeDocuments + ) { Queue queue = new LinkedList<>(); queue.offer(new Object[] { root, false }); while (!queue.isEmpty()) { Object[] current = queue.poll(); - TreeNode node = (TreeNode) current[0]; + InMemoryTreeNode node = (InMemoryTreeNode) current[0]; boolean currentIsStarNode = (boolean) current[1]; assertNotNull(node); // assert dimensions - if (node.dimensionId != TreeNode.ALL) { + if (node.dimensionId != StarTreeUtils.ALL) { assertTrue(node.dimensionId >= 0 && node.dimensionId < totalDimensions); } if (node.children != null && !node.children.isEmpty()) { assertEquals(node.dimensionId + 1, node.childDimensionId); assertTrue(node.childDimensionId < totalDimensions); - TreeNode starNode = null; + InMemoryTreeNode starNode = null; Object[] nonStarNodeCumulativeMetrics = getMetrics(starTreeDocuments); - for (Map.Entry entry : node.children.entrySet()) { + for (Map.Entry entry : node.children.entrySet()) { Long childDimensionValue = entry.getKey(); - TreeNode child = entry.getValue(); + InMemoryTreeNode child = entry.getValue(); Object[] currMetrics = getMetrics(starTreeDocuments); - if (!child.isStarNode) { + if (child.nodeType != StarTreeNodeType.STAR.getValue()) { // Validate dimension values in documents for (int i = child.startDocId; i < child.endDocId; i++) { StarTreeDocument doc = starTreeDocuments.get(i); int j = 0; addMetrics(doc, currMetrics, j); - if (!child.isStarNode) { + if (child.nodeType != StarTreeNodeType.STAR.getValue()) { Long dimension = doc.dimensions[child.dimensionId]; assertEquals(childDimensionValue, dimension); if (dimension != null) { assertEquals(child.dimensionValue, (long) dimension); } else { // TODO : fix this ? - assertEquals(child.dimensionValue, TreeNode.ALL); + assertEquals(child.dimensionValue, StarTreeUtils.ALL); } } } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeFileFormatsTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeFileFormatsTests.java index 5f7a48f8f731c..d638f697e8441 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeFileFormatsTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/data/StarTreeFileFormatsTests.java @@ -82,10 +82,10 @@ public void test_StarTreeNode() throws IOException { StarTreeNode child = childrenIterator.next(); if (child.getStarTreeNodeType() == StarTreeNodeType.DEFAULT.getValue()) { assertStarTreeNode( - starTreeNode.getChildForDimensionValue(child.getDimensionValue(), false), + starTreeNode.getChildForDimensionValue(child.getDimensionValue()), inMemoryTreeNodeMap.get(child.getDimensionValue()) ); - assertNull(starTreeNode.getChildForDimensionValue(child.getDimensionValue(), true)); + assertNull(starTreeNode.getChildStarNode()); } queue.add(child); @@ -124,7 +124,7 @@ public void test_starTreeSearch() throws IOException { for (int i = 0; i < maxLevels - 1; i++) { InMemoryTreeNode randomChildNode = randomFrom(inMemoryTreeNode.children.values()); - StarTreeNode randomStarTreeChildNode = starTreeNode.getChildForDimensionValue(randomChildNode.dimensionValue, false); + StarTreeNode randomStarTreeChildNode = starTreeNode.getChildForDimensionValue(randomChildNode.dimensionValue); assertNotNull(randomStarTreeChildNode); assertStarTreeNode(randomStarTreeChildNode, randomChildNode); diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaTests.java index b3845e2d31e2f..74d9244e16010 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetaTests.java @@ -182,10 +182,19 @@ public void test_starTreeMetadata() throws IOException { assertEquals(dimensionsOrder.get(i).getField(), starTreeMetadata.getDimensionFields().get(i)); } - for (int i = 0; i < metricAggregatorInfos.size(); i++) { - MetricEntry metricEntry = starTreeMetadata.getMetricEntries().get(i); - assertEquals(metricAggregatorInfos.get(i).getField(), metricEntry.getMetricFieldName()); - assertEquals(metricAggregatorInfos.get(i).getMetricStat(), metricEntry.getMetricStat()); + assertEquals(starTreeField.getMetrics().size(), starTreeMetadata.getMetrics().size()); + + for (int i = 0; i < starTreeField.getMetrics().size(); i++) { + + Metric expectedMetric = starTreeField.getMetrics().get(i); + Metric resultMetric = starTreeMetadata.getMetrics().get(i); + + assertEquals(expectedMetric.getField(), resultMetric.getField()); + assertEquals(expectedMetric.getMetrics().size(), resultMetric.getMetrics().size()); + + for (int j = 0; j < resultMetric.getMetrics().size(); j++) { + assertEquals(expectedMetric.getMetrics().get(j), resultMetric.getMetrics().get(j)); + } } assertEquals(segmentDocumentCount, starTreeMetadata.getSegmentAggregatedDocCount(), 0); assertEquals(maxLeafDocs, starTreeMetadata.getMaxLeafDocs(), 0); diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNodeTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java similarity index 93% rename from server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNodeTests.java rename to server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java index 61163a37582ad..daf2acbcf2365 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/node/FixedLengthStarTreeNodeTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.index.compositeindex.datacube.startree.node; +package org.opensearch.index.compositeindex.datacube.startree.fileformats.node; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; @@ -14,6 +14,8 @@ import org.apache.lucene.store.IndexOutput; import org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter; import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTree; import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; import org.opensearch.test.OpenSearchTestCase; import org.junit.Before; @@ -139,7 +141,7 @@ public void testGetStarTreeNodeType() throws IOException { public void testGetChildForDimensionValue() throws IOException { long dimensionValue = randomIntBetween(0, node.children.size() - 2); - FixedLengthStarTreeNode childNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue(dimensionValue, false); + FixedLengthStarTreeNode childNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue(dimensionValue); assertNotNull(childNode); assertEquals(dimensionValue, childNode.getDimensionValue()); } @@ -157,19 +159,19 @@ public void testGetChildrenIterator() throws IOException { public void testGetChildForStarNode() throws IOException { // Assuming the first child is a star node in our test data - FixedLengthStarTreeNode starNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue((long) StarTreeUtils.ALL, true); + FixedLengthStarTreeNode starNode = (FixedLengthStarTreeNode) starTreeNode.getChildStarNode(); assertNotNull(starNode); assertEquals(StarTreeUtils.ALL, starNode.getDimensionValue()); } public void testGetChildForNullNode() throws IOException { - FixedLengthStarTreeNode nullNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue(null, false); + FixedLengthStarTreeNode nullNode = (FixedLengthStarTreeNode) starTreeNode.getChildForDimensionValue(null); assertNull(nullNode); } public void testGetChildForInvalidDimensionValue() throws IOException { long invalidDimensionValue = Long.MAX_VALUE; - assertThrows(AssertionError.class, () -> starTreeNode.getChildForDimensionValue(invalidDimensionValue, false)); + assertThrows(AssertionError.class, () -> starTreeNode.getChildForDimensionValue(invalidDimensionValue)); } public void tearDown() throws Exception {