From 19b39249d607ddd40f08b86de309ebf3bc12d4d1 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Mon, 10 Jun 2024 18:37:12 +0530 Subject: [PATCH] Refactoring Signed-off-by: Sarthak Aggarwal --- ...er.java => BaseSingleStarTreeBuilder.java} | 148 ++++++------- .../aggregators/MetricTypeFieldPair.java | 40 ++-- .../aggregators/SumValueAggregator.java | 1 + .../startree/aggregators/ValueAggregator.java | 1 + .../aggregators/ValueAggregatorFactory.java | 1 + ...ieldWriter.java => SingleTreeBuilder.java} | 4 +- .../{aggregators => data}/DataType.java | 2 +- .../{builder => data}/StarTreeDocValues.java | 2 +- .../startree/data/StarTreeDocument.java | 29 +++ .../startree/node/StarTreeNode.java | 1 + .../startree/utils/StarTreeBuilderUtils.java | 16 +- .../index/TestBaseSingleStarTreeBuilder.java | 203 ++++++++++++++++++ .../aggregators/MetricTypeFieldPairTests.java | 8 +- .../ValueAggregatorFactoryTests.java | 1 + 14 files changed, 339 insertions(+), 118 deletions(-) rename server/src/main/java/org/apache/lucene/index/{BaseCompositeFieldStarTreeBuilder.java => BaseSingleStarTreeBuilder.java} (87%) rename server/src/main/java/org/opensearch/index/compositeindex/startree/builder/{CompositeFieldWriter.java => SingleTreeBuilder.java} (84%) rename server/src/main/java/org/opensearch/index/compositeindex/startree/{aggregators => data}/DataType.java (96%) rename server/src/main/java/org/opensearch/index/compositeindex/startree/{builder => data}/StarTreeDocValues.java (94%) create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocument.java create mode 100644 server/src/test/java/org/apache/lucene/index/TestBaseSingleStarTreeBuilder.java diff --git a/server/src/main/java/org/apache/lucene/index/BaseCompositeFieldStarTreeBuilder.java b/server/src/main/java/org/apache/lucene/index/BaseSingleStarTreeBuilder.java similarity index 87% rename from server/src/main/java/org/apache/lucene/index/BaseCompositeFieldStarTreeBuilder.java rename to server/src/main/java/org/apache/lucene/index/BaseSingleStarTreeBuilder.java index 4c78976c0d4a7..94494e355c087 100644 --- a/server/src/main/java/org/apache/lucene/index/BaseCompositeFieldStarTreeBuilder.java +++ b/server/src/main/java/org/apache/lucene/index/BaseSingleStarTreeBuilder.java @@ -9,15 +9,12 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.ByteBlockPool; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.Counter; -import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.compositeindex.CompositeField; import org.opensearch.index.compositeindex.DateDimension; import org.opensearch.index.compositeindex.Dimension; @@ -27,9 +24,10 @@ import org.opensearch.index.compositeindex.startree.aggregators.MetricTypeFieldPair; import org.opensearch.index.compositeindex.startree.aggregators.ValueAggregator; import org.opensearch.index.compositeindex.startree.aggregators.ValueAggregatorFactory; -import org.opensearch.index.compositeindex.startree.builder.CompositeFieldWriter; -import org.opensearch.index.compositeindex.startree.builder.StarTreeDocValues; +import org.opensearch.index.compositeindex.startree.builder.SingleTreeBuilder; import org.opensearch.index.compositeindex.startree.builder.StarTreeDocValuesIteratorFactory; +import org.opensearch.index.compositeindex.startree.data.StarTreeDocValues; +import org.opensearch.index.compositeindex.startree.data.StarTreeDocument; import org.opensearch.index.compositeindex.startree.node.StarTreeNode; import org.opensearch.index.compositeindex.startree.utils.StarTreeBuilderUtils; import org.opensearch.index.mapper.NumberFieldMapper; @@ -47,35 +45,36 @@ /** * Base class for star tree builder */ -public abstract class BaseCompositeFieldStarTreeBuilder implements CompositeFieldWriter { +public abstract class BaseSingleStarTreeBuilder implements SingleTreeBuilder { // TODO: STAR_TREE_CODEC will be moved to CodecService once the Star Tree Codec is defined public static final String STAR_TREE_CODEC = "startreecodec"; - private static final Logger logger = LogManager.getLogger(BaseCompositeFieldStarTreeBuilder.class); + private static final Logger logger = LogManager.getLogger(BaseSingleStarTreeBuilder.class); public static final int STAR_IN_DOC_VALUES_INDEX = 0; - public final String[] dimensionsSplitOrder; - public final Set skipStarNodeCreationForDimensions; - public final String[] metrics; + protected final String[] dimensionsSplitOrder; + protected final Set skipStarNodeCreationForDimensions; + protected final String[] metrics; - public final int numMetrics; - public final int numDimensions; - public int numDocs; - public int totalDocs; - public int numNodes; - public final int maxLeafDocuments; + protected final int numMetrics; + protected final int numDimensions; + protected int numDocs; + protected int totalDocs; + protected int numNodes; + protected final int maxLeafDocuments; - public final StarTreeBuilderUtils.TreeNode rootNode = getNewNode(); + protected final StarTreeBuilderUtils.TreeNode rootNode = getNewNode(); - public IndexOutput indexOutput; - public DocIdSetIterator[] dimensionReaders; - public DocIdSetIterator[] metricReaders; + // TODO: This will be initialized with OnHeap / OffHeap Implementations (Commented it's occurrences for now) + // private IndexOutput indexOutput; + protected DocIdSetIterator[] dimensionReaders; + protected DocIdSetIterator[] metricReaders; - public ValueAggregator[] valueAggregators; - public DocValuesConsumer docValuesConsumer; - public DocValuesProducer docValuesProducer; + protected ValueAggregator[] valueAggregators; + protected DocValuesConsumer docValuesConsumer; + protected DocValuesProducer docValuesProducer; private final StarTreeDocValuesIteratorFactory starTreeDocValuesIteratorFactory; private final CompositeField compositeField; @@ -90,7 +89,7 @@ public abstract class BaseCompositeFieldStarTreeBuilder implements CompositeFiel * @param docValuesConsumer to consume the new aggregated metrics during flush * @param state stores the segment state */ - protected BaseCompositeFieldStarTreeBuilder( + protected BaseSingleStarTreeBuilder( CompositeField compositeField, DocValuesProducer docValuesProducer, DocValuesConsumer docValuesConsumer, @@ -99,11 +98,11 @@ protected BaseCompositeFieldStarTreeBuilder( logger.info("Building in base star tree builder"); - String docFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "stttree"); - logger.info("Star tree file name : {}", docFileName); + // String docFileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, "stttree"); + // logger.info("Star tree file name : {}", docFileName); - indexOutput = state.directory.createOutput(docFileName, state.context); - CodecUtil.writeIndexHeader(indexOutput, STAR_TREE_CODEC, 0, state.segmentInfo.getId(), state.segmentSuffix); + // indexOutput = state.directory.createOutput(docFileName, state.context); + // CodecUtil.writeIndexHeader(indexOutput, STAR_TREE_CODEC, 0, state.segmentInfo.getId(), state.segmentSuffix); starTreeDocValuesIteratorFactory = new StarTreeDocValuesIteratorFactory(); this.compositeField = compositeField; @@ -145,7 +144,7 @@ protected BaseCompositeFieldStarTreeBuilder( int index = 0; for (MetricTypeFieldPair metricTypeFieldPair : metricTypeFieldPairs) { metrics[index] = metricTypeFieldPair.toFieldName(); - valueAggregators[index] = ValueAggregatorFactory.getValueAggregator(metricTypeFieldPair.getFunctionType()); + valueAggregators[index] = ValueAggregatorFactory.getValueAggregator(metricTypeFieldPair.getMetricType()); // Ignore the column for COUNT aggregation function if (valueAggregators[index].getAggregationType() != MetricType.COUNT) { String metricName = metricTypeFieldPair.getField(); @@ -165,7 +164,7 @@ protected BaseCompositeFieldStarTreeBuilder( /** * Generates the AggregationFunctionColumnPairs for all the metrics on a field */ - private List generateAggregationFunctionColumnPairs() { + List generateAggregationFunctionColumnPairs() { List metricTypeFieldPairs = new ArrayList<>(); for (Metric metric : this.compositeField.getMetrics()) { for (MetricType metricType : metric.getMetrics()) { @@ -293,17 +292,17 @@ protected StarTreeDocument mergeSegmentStarTreeDocument( ) { // TODO: HANDLE KEYWORDS LATER! if (aggregatedStarTreeDocument == null) { - long[] dimensions = Arrays.copyOf(segmentStarTreeDocument._dimensions, numDimensions); + long[] dimensions = Arrays.copyOf(segmentStarTreeDocument.dimensions, numDimensions); Object[] metrics = new Object[numMetrics]; for (int i = 0; i < numMetrics; i++) { - metrics[i] = valueAggregators[i].getInitialAggregatedValue(segmentStarTreeDocument._metrics[i]); + metrics[i] = valueAggregators[i].getInitialAggregatedValue(segmentStarTreeDocument.metrics[i]); } return new StarTreeDocument(dimensions, metrics); } else { for (int i = 0; i < numMetrics; i++) { - aggregatedStarTreeDocument._metrics[i] = valueAggregators[i].applyRawValue( - aggregatedStarTreeDocument._metrics[i], - segmentStarTreeDocument._metrics[i] + aggregatedStarTreeDocument.metrics[i] = valueAggregators[i].applyRawValue( + aggregatedStarTreeDocument.metrics[i], + segmentStarTreeDocument.metrics[i] ); } return aggregatedStarTreeDocument; @@ -324,17 +323,17 @@ protected StarTreeDocument mergeStarTreeDocument( StarTreeDocument starTreeStarTreeDocument ) { if (aggregatedStarTreeDocument == null) { - long[] dimensions = Arrays.copyOf(starTreeStarTreeDocument._dimensions, numDimensions); + long[] dimensions = Arrays.copyOf(starTreeStarTreeDocument.dimensions, numDimensions); Object[] metrics = new Object[numMetrics]; for (int i = 0; i < numMetrics; i++) { - metrics[i] = valueAggregators[i].cloneAggregatedValue((Long) starTreeStarTreeDocument._metrics[i]); + metrics[i] = valueAggregators[i].cloneAggregatedValue(starTreeStarTreeDocument.metrics[i]); } return new StarTreeDocument(dimensions, metrics); } else { for (int i = 0; i < numMetrics; i++) { - aggregatedStarTreeDocument._metrics[i] = valueAggregators[i].applyAggregatedValue( - (Long) starTreeStarTreeDocument._metrics[i], - (Long) aggregatedStarTreeDocument._metrics[i] + aggregatedStarTreeDocument.metrics[i] = valueAggregators[i].applyAggregatedValue( + starTreeStarTreeDocument.metrics[i], + aggregatedStarTreeDocument.metrics[i] ); } return aggregatedStarTreeDocument; @@ -366,7 +365,7 @@ public void build(Iterator starTreeDocumentIterator) throws IO logger.info("Generated star tree docs : [{}] from segment docs : [{}]", numStarTreeDocument, numSegmentStarTreeDocument); if (numDocs == 0) { - StarTreeBuilderUtils.serializeTree(indexOutput, rootNode, dimensionsSplitOrder, numNodes); + // StarTreeBuilderUtils.serializeTree(indexOutput, rootNode, dimensionsSplitOrder, numNodes); return; } @@ -386,7 +385,7 @@ public void build(Iterator starTreeDocumentIterator) throws IO createSortedDocValuesIndices(docValuesConsumer); // Serialize and save in disk - StarTreeBuilderUtils.serializeTree(indexOutput, rootNode, dimensionsSplitOrder, numNodes); + // StarTreeBuilderUtils.serializeTree(indexOutput, rootNode, dimensionsSplitOrder, numNodes); // TODO: Write star tree metadata for off heap implementation @@ -502,7 +501,7 @@ private StarTreeDocument createAggregatedDocs(StarTreeBuilderUtils.TreeNode node } assert aggregatedStarTreeDocument != null; for (int i = node.dimensionId + 1; i < numDimensions; i++) { - aggregatedStarTreeDocument._dimensions[i] = STAR_IN_DOC_VALUES_INDEX; + aggregatedStarTreeDocument.dimensions[i] = STAR_IN_DOC_VALUES_INDEX; } node.aggregatedDocId = numDocs; appendToStarTree(aggregatedStarTreeDocument); @@ -526,7 +525,7 @@ private StarTreeDocument createAggregatedDocs(StarTreeBuilderUtils.TreeNode node } assert aggregatedStarTreeDocument != null; for (int i = node.dimensionId + 1; i < numDimensions; i++) { - aggregatedStarTreeDocument._dimensions[i] = STAR_IN_DOC_VALUES_INDEX; + aggregatedStarTreeDocument.dimensions[i] = STAR_IN_DOC_VALUES_INDEX; } node.aggregatedDocId = numDocs; appendToStarTree(aggregatedStarTreeDocument); @@ -608,8 +607,8 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer) t Map> ordinalToSortedSetDocValueMap = new HashMap<>(); for (int docId = 0; docId < numDocs; docId++) { StarTreeDocument starTreeDocument = getStarTreeDocument(docId); - for (int i = 0; i < starTreeDocument._dimensions.length; i++) { - long val = starTreeDocument._dimensions[i]; + for (int i = 0; i < starTreeDocument.dimensions.length; i++) { + long val = starTreeDocument.dimensions[i]; StarTreeDocValuesWriter starTreeDocValuesWriter = dimensionWriters.get(i); switch (starTreeDocValuesWriter.getDocValuesType()) { case SORTED_SET: @@ -626,9 +625,9 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer) t throw new IllegalStateException("Unsupported doc values type"); } } - for (int i = 0; i < starTreeDocument._metrics.length; i++) { + for (int i = 0; i < starTreeDocument.metrics.length; i++) { try { - Number parse = NumberFieldMapper.NumberType.LONG.parse(starTreeDocument._metrics[i], true); + Number parse = NumberFieldMapper.NumberType.LONG.parse(starTreeDocument.metrics[i], true); StarTreeDocValuesWriter starTreeDocValuesWriter = metricWriters.get(i); ((SortedNumericDocValuesWriter) starTreeDocValuesWriter.getDocValuesWriter()).addValue(docId, parse.longValue()); } catch (IllegalArgumentException e) { @@ -686,7 +685,7 @@ public SortedSetDocValues getSortedSet(FieldInfo field) { /** * Returns the respective doc values writer based on doc values type */ - private DocValuesWriter getDocValuesWriter(DocValuesType docValuesType, FieldInfo fi, Counter counter) { + DocValuesWriter getDocValuesWriter(DocValuesType docValuesType, FieldInfo fi, Counter counter) { final ByteBlockPool.Allocator byteBlockAllocator = new ByteBlockPool.DirectTrackingAllocator(counter); final ByteBlockPool docValuesBytePool = new ByteBlockPool(byteBlockAllocator); switch (docValuesType) { @@ -708,40 +707,23 @@ private long handleDateDimension(final String fieldName, final long val) { } public void close() throws IOException { - boolean success = false; - try { - if (indexOutput != null) { - indexOutput.writeInt(-1); - CodecUtil.writeFooter(indexOutput); // write checksum - } - success = true; - } catch (Exception e) { - throw new RuntimeException(e); - } finally { - if (success) { - IOUtils.close(indexOutput); - } else { - IOUtils.closeWhileHandlingException(indexOutput); - } - indexOutput = null; - } + // boolean success = false; + // try { + // if (indexOutput != null) { + // indexOutput.writeInt(-1); + // CodecUtil.writeFooter(indexOutput); // write checksum + // } + // success = true; + // } catch (Exception e) { + // throw new RuntimeException(e); + // } finally { + // if (success) { + // IOUtils.close(indexOutput); + // } else { + // IOUtils.closeWhileHandlingException(indexOutput); + // } + // indexOutput = null; + // } } - /** - * Star tree document - */ - public static class StarTreeDocument { - public final long[] _dimensions; - public final Object[] _metrics; - - public StarTreeDocument(long[] dimensions, Object[] metrics) { - _dimensions = dimensions; - _metrics = metrics; - } - - @Override - public String toString() { - return Arrays.toString(_dimensions) + " | " + Arrays.toString(_metrics); - } - } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPair.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPair.java index f5e39c0610119..ca7caaa5535ee 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPair.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPair.java @@ -21,15 +21,15 @@ public class MetricTypeFieldPair implements Comparable { public static final String STAR = "*"; public static final MetricTypeFieldPair COUNT_STAR = new MetricTypeFieldPair(MetricType.COUNT, STAR); - private final MetricType functionType; + private final MetricType metricType; private final String field; /** * Constructor for MetricTypeFieldPair */ - public MetricTypeFieldPair(MetricType functionType, String field) { - this.functionType = functionType; - if (functionType == MetricType.COUNT) { + public MetricTypeFieldPair(MetricType metricType, String field) { + this.metricType = metricType; + if (metricType == MetricType.COUNT) { this.field = STAR; } else { this.field = field; @@ -39,8 +39,8 @@ public MetricTypeFieldPair(MetricType functionType, String field) { /** * @return Metric Type */ - public MetricType getFunctionType() { - return functionType; + public MetricType getMetricType() { + return metricType; } /** @@ -51,17 +51,17 @@ public String getField() { } /** - * @return field name with function type and field + * @return field name with metric type and field */ public String toFieldName() { - return toFieldName(functionType, field); + return toFieldName(metricType, field); } /** - * Builds field name with function type and field + * Builds field name with metric type and field */ - public static String toFieldName(MetricType functionType, String field) { - return functionType.getTypeName() + DELIMITER + field; + public static String toFieldName(MetricType metricType, String field) { + return metricType.getTypeName() + DELIMITER + field; } /** @@ -69,24 +69,24 @@ public static String toFieldName(MetricType functionType, String field) { */ public static MetricTypeFieldPair fromFieldName(String fieldName) { String[] parts = fieldName.split(DELIMITER, 2); - return fromFunctionAndFieldName(parts[0], parts[1]); + return fromMetricAndFieldName(parts[0], parts[1]); } /** - * Builds MetricTypeFieldPair from function and field name + * Builds MetricTypeFieldPair from metric and field name */ - private static MetricTypeFieldPair fromFunctionAndFieldName(String functionName, String fieldName) { - MetricType functionType = MetricType.fromTypeName(functionName); - if (functionType == MetricType.COUNT) { + private static MetricTypeFieldPair fromMetricAndFieldName(String metricName, String fieldName) { + MetricType metricType = MetricType.fromTypeName(metricName); + if (metricType == MetricType.COUNT) { return COUNT_STAR; } else { - return new MetricTypeFieldPair(functionType, fieldName); + return new MetricTypeFieldPair(metricType, fieldName); } } @Override public int hashCode() { - return 31 * functionType.hashCode() + field.hashCode(); + return 31 * metricType.hashCode() + field.hashCode(); } @Override @@ -96,7 +96,7 @@ public boolean equals(Object obj) { } if (obj instanceof MetricTypeFieldPair) { MetricTypeFieldPair anotherPair = (MetricTypeFieldPair) obj; - return functionType == anotherPair.functionType && field.equals(anotherPair.field); + return metricType == anotherPair.metricType && field.equals(anotherPair.field); } return false; } @@ -109,7 +109,7 @@ public String toString() { @Override public int compareTo(MetricTypeFieldPair other) { return Comparator.comparing((MetricTypeFieldPair o) -> o.field) - .thenComparing((MetricTypeFieldPair o) -> o.functionType) + .thenComparing((MetricTypeFieldPair o) -> o.metricType) .compare(this, other); } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/SumValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/SumValueAggregator.java index c3bf11f1775fa..9fbfdd26df57c 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/SumValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/SumValueAggregator.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.startree.aggregators; import org.opensearch.index.compositeindex.MetricType; +import org.opensearch.index.compositeindex.startree.data.DataType; /** * Sum value aggregator for star tree diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregator.java index cc3410f3e3e7e..bfbbc05254d5e 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregator.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.startree.aggregators; import org.opensearch.index.compositeindex.MetricType; +import org.opensearch.index.compositeindex.startree.data.DataType; /** * A value aggregator that pre-aggregates on the input values for a specific type of aggregation. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactory.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactory.java index 20a054b30f64d..bb1df58ca6b2b 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactory.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactory.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.startree.aggregators; import org.opensearch.index.compositeindex.MetricType; +import org.opensearch.index.compositeindex.startree.data.DataType; /** * Value aggregator factory for a given aggregation type diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/CompositeFieldWriter.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/SingleTreeBuilder.java similarity index 84% rename from server/src/main/java/org/opensearch/index/compositeindex/startree/builder/CompositeFieldWriter.java rename to server/src/main/java/org/opensearch/index/compositeindex/startree/builder/SingleTreeBuilder.java index 84a197fdcf1df..2d485db89b474 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/CompositeFieldWriter.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/SingleTreeBuilder.java @@ -8,6 +8,8 @@ package org.opensearch.index.compositeindex.startree.builder; +import org.opensearch.index.compositeindex.startree.data.StarTreeDocValues; + import java.io.Closeable; import java.io.IOException; import java.util.List; @@ -16,7 +18,7 @@ * A star-tree builder that builds a single star-tree. * @opensearch.experimental */ -public interface CompositeFieldWriter extends Closeable { +public interface SingleTreeBuilder extends Closeable { /** * Builds the data structure for the given composite index config. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/DataType.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/DataType.java similarity index 96% rename from server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/DataType.java rename to server/src/main/java/org/opensearch/index/compositeindex/startree/data/DataType.java index 5dc04d4196659..c15697d722563 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/aggregators/DataType.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/DataType.java @@ -5,7 +5,7 @@ * this file be licensed under the Apache-2.0 license or a * compatible open source license. */ -package org.opensearch.index.compositeindex.startree.aggregators; +package org.opensearch.index.compositeindex.startree.data; /** * Data type of doc values diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/StarTreeDocValues.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocValues.java similarity index 94% rename from server/src/main/java/org/opensearch/index/compositeindex/startree/builder/StarTreeDocValues.java rename to server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocValues.java index abebd21444cc4..bdedcefb97a61 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/builder/StarTreeDocValues.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocValues.java @@ -5,7 +5,7 @@ * this file be licensed under the Apache-2.0 license or a * compatible open source license. */ -package org.opensearch.index.compositeindex.startree.builder; +package org.opensearch.index.compositeindex.startree.data; import org.apache.lucene.index.SortedNumericDocValues; import org.opensearch.index.compositeindex.startree.node.StarTree; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocument.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocument.java new file mode 100644 index 0000000000000..44b211123c19f --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/data/StarTreeDocument.java @@ -0,0 +1,29 @@ +/* + * 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.startree.data; + +import java.util.Arrays; + +/** + * Star tree document + */ +public class StarTreeDocument { + public final long[] dimensions; + public final Object[] metrics; + + public StarTreeDocument(long[] dimensions, Object[] metrics) { + this.dimensions = dimensions; + this.metrics = metrics; + } + + @Override + public String toString() { + return Arrays.toString(dimensions) + " | " + Arrays.toString(metrics); + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/node/StarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/node/StarTreeNode.java index c3ad5406a2872..ea33f817a689b 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/node/StarTreeNode.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/node/StarTreeNode.java @@ -12,6 +12,7 @@ /** * Class representing each node in star tree + * The interface is implemented by build-mode based nodes * @opensearch.experimental */ public interface StarTreeNode { diff --git a/server/src/main/java/org/opensearch/index/compositeindex/startree/utils/StarTreeBuilderUtils.java b/server/src/main/java/org/opensearch/index/compositeindex/startree/utils/StarTreeBuilderUtils.java index 5f86492149c23..5c6b1dcdda149 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/startree/utils/StarTreeBuilderUtils.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/startree/utils/StarTreeBuilderUtils.java @@ -37,17 +37,17 @@ public class StarTreeBuilderUtils { private StarTreeBuilderUtils() {} - public static final int INVALID_ID = -1; + public static final int ALL = -1; public static final long MAGIC_MARKER = 0xBADDA55B00DAD00DL; /** Tree node representation */ public static class TreeNode { - public int dimensionId = INVALID_ID; - public long dimensionValue = INVALID_ID; - public int startDocId = INVALID_ID; - public int endDocId = INVALID_ID; - public int aggregatedDocId = INVALID_ID; - public int childDimensionId = INVALID_ID; + public int dimensionId = ALL; + public long dimensionValue = ALL; + public int startDocId = ALL; + public int endDocId = ALL; + public int aggregatedDocId = ALL; + public int childDimensionId = ALL; public Map children; } @@ -100,7 +100,7 @@ static void writeNodes(IndexOutput output, TreeNode rootNode) throws IOException TreeNode node = queue.remove(); if (node.children == null) { - writeNode(output, node, INVALID_ID, INVALID_ID); + writeNode(output, node, ALL, ALL); } else { // Sort all children nodes based on dimension value List sortedChildren = new ArrayList<>(node.children.values()); diff --git a/server/src/test/java/org/apache/lucene/index/TestBaseSingleStarTreeBuilder.java b/server/src/test/java/org/apache/lucene/index/TestBaseSingleStarTreeBuilder.java new file mode 100644 index 0000000000000..471319f304957 --- /dev/null +++ b/server/src/test/java/org/apache/lucene/index/TestBaseSingleStarTreeBuilder.java @@ -0,0 +1,203 @@ +/* + * 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.apache.lucene.index; + +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.Counter; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.Version; +import org.opensearch.index.compositeindex.CompositeField; +import org.opensearch.index.compositeindex.Dimension; +import org.opensearch.index.compositeindex.Metric; +import org.opensearch.index.compositeindex.MetricType; +import org.opensearch.index.compositeindex.StarTreeFieldSpec; +import org.opensearch.index.compositeindex.startree.aggregators.MetricTypeFieldPair; +import org.opensearch.index.compositeindex.startree.data.StarTreeDocValues; +import org.opensearch.index.compositeindex.startree.data.StarTreeDocument; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Set; +import java.util.UUID; + +import static org.mockito.Mockito.mock; + +public class TestBaseSingleStarTreeBuilder extends OpenSearchTestCase { + + private static BaseSingleStarTreeBuilder builder; + private static List dimensionsOrder; + private static List fields = List.of( + "field1", + "field2", + "field3", + "field4", + "field5", + "field6", + "field7", + "field8", + "field9", + "field10" + ); + private static List metrics; + private static Directory directory; + private static FieldInfo[] fieldsInfo; + + @BeforeClass + public static void setup() throws IOException { + + dimensionsOrder = List.of(new Dimension("field1"), new Dimension("field3"), new Dimension("field5"), new Dimension("field8")); + metrics = List.of(new Metric("field2", List.of(MetricType.SUM)), new Metric("field4", List.of(MetricType.SUM))); + + CompositeField compositeField = new CompositeField( + "test", + dimensionsOrder, + metrics, + new StarTreeFieldSpec(1, Set.of("field8"), StarTreeFieldSpec.StarTreeBuildMode.ON_HEAP) + ); + DocValuesConsumer docValuesConsumer = mock(DocValuesConsumer.class); + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + directory = newFSDirectory(createTempDir()); + SegmentInfo segmentInfo = new SegmentInfo( + directory, + Version.LATEST, + Version.LUCENE_9_11_0, + "test_segment", + 5, + false, + false, + new Lucene99Codec(), + new HashMap<>(), + UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), + new HashMap<>(), + null + ); + + fieldsInfo = new FieldInfo[fields.size()]; + + for (int i = 0; i < fieldsInfo.length; i++) { + fieldsInfo[i] = new FieldInfo( + fields.get(i), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + } + FieldInfos fieldInfos = new FieldInfos(fieldsInfo); + final SegmentWriteState state = new SegmentWriteState( + InfoStream.getDefault(), + segmentInfo.dir, + segmentInfo, + fieldInfos, + null, + newIOContext(random()) + ); + + builder = new BaseSingleStarTreeBuilder(compositeField, docValuesProducer, docValuesConsumer, state) { + @Override + public void appendStarTreeDocument(StarTreeDocument starTreeDocument) throws IOException {} + + @Override + public StarTreeDocument getStarTreeDocument(int docId) throws IOException { + return null; + } + + @Override + public long getDimensionValue(int docId, int dimensionId) throws IOException { + return 0; + } + + @Override + public Iterator sortAndAggregateSegmentStarTreeDocument(int numDocs) throws IOException { + return null; + } + + @Override + public Iterator generateStarTreeForStarNode(int startDocId, int endDocId, int dimensionId) + throws IOException { + return null; + } + + @Override + public void build(List starTreeDocValues) throws IOException { + + } + }; + } + + public void test_generateAggregationFunctionColumnPairs() throws IOException { + List metricTypeFieldPairs = builder.generateAggregationFunctionColumnPairs(); + List expectedMetricTypeFieldPairs = List.of( + new MetricTypeFieldPair(MetricType.SUM, "field2"), + new MetricTypeFieldPair(MetricType.SUM, "field4") + ); + assertEquals(metricTypeFieldPairs, expectedMetricTypeFieldPairs); + } + + public void test_mergeStarTreeDocument() { + StarTreeDocument starTreeDocument1 = new StarTreeDocument(new long[] { 1, 3, 5, 8 }, new Double[] { 4.0, 8.0 }); + StarTreeDocument starTreeDocument2 = new StarTreeDocument(new long[] { 1, 3, 5, 8 }, new Double[] { 10.0, 6.0 }); + + StarTreeDocument expectedeMergedStarTreeDocument = new StarTreeDocument(new long[] { 1, 3, 5, 8 }, new Double[] { 14.0, 14.0 }); + StarTreeDocument mergedStarTreeDocument = builder.mergeStarTreeDocument(starTreeDocument1, starTreeDocument2); + + assertEquals(mergedStarTreeDocument.metrics[0], expectedeMergedStarTreeDocument.metrics[0]); + assertEquals(mergedStarTreeDocument.metrics[1], expectedeMergedStarTreeDocument.metrics[1]); + } + + public void test_mergeStarTreeDocument_nullAggregatedStarTreeDocument() { + StarTreeDocument starTreeDocument = new StarTreeDocument(new long[] { 1, 3, 5, 8 }, new Double[] { 10.0, 6.0 }); + + StarTreeDocument expectedeMergedStarTreeDocument = new StarTreeDocument(new long[] { 1, 3, 5, 8 }, new Double[] { 10.0, 6.0 }); + StarTreeDocument mergedStarTreeDocument = builder.mergeStarTreeDocument(null, starTreeDocument); + + assertEquals(mergedStarTreeDocument.metrics[0], expectedeMergedStarTreeDocument.metrics[0]); + assertEquals(mergedStarTreeDocument.metrics[1], expectedeMergedStarTreeDocument.metrics[1]); + } + + public void test_getDocValuesWriter() { + assertTrue( + builder.getDocValuesWriter(DocValuesType.SORTED_SET, fieldsInfo[0], Counter.newCounter()) instanceof SortedSetDocValuesWriter + ); + assertTrue( + builder.getDocValuesWriter( + DocValuesType.SORTED_NUMERIC, + fieldsInfo[0], + Counter.newCounter() + ) instanceof SortedNumericDocValuesWriter + ); + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + directory.close(); + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPairTests.java b/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPairTests.java index 0827d86162754..1dca82708a19c 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPairTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/MetricTypeFieldPairTests.java @@ -15,13 +15,13 @@ public class MetricTypeFieldPairTests extends OpenSearchTestCase { public void testConstructor() { MetricTypeFieldPair pair = new MetricTypeFieldPair(MetricType.SUM, "column1"); - assertEquals(MetricType.SUM, pair.getFunctionType()); + assertEquals(MetricType.SUM, pair.getMetricType()); assertEquals("column1", pair.getField()); } public void testCountStarConstructor() { MetricTypeFieldPair pair = new MetricTypeFieldPair(MetricType.COUNT, "anything"); - assertEquals(MetricType.COUNT, pair.getFunctionType()); + assertEquals(MetricType.COUNT, pair.getMetricType()); assertEquals("*", pair.getField()); } @@ -32,13 +32,13 @@ public void testToFieldName() { public void testFromFieldName() { MetricTypeFieldPair pair = MetricTypeFieldPair.fromFieldName("max__column3"); - assertEquals(MetricType.MAX, pair.getFunctionType()); + assertEquals(MetricType.MAX, pair.getMetricType()); assertEquals("column3", pair.getField()); } public void testCountStarFromFieldName() { MetricTypeFieldPair pair = MetricTypeFieldPair.fromFieldName("count__*"); - assertEquals(MetricType.COUNT, pair.getFunctionType()); + assertEquals(MetricType.COUNT, pair.getMetricType()); assertEquals("*", pair.getField()); assertSame(MetricTypeFieldPair.COUNT_STAR, pair); } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactoryTests.java b/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactoryTests.java index 60129e0a7dc96..3ed8d0a8858e9 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactoryTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/startree/aggregators/ValueAggregatorFactoryTests.java @@ -9,6 +9,7 @@ package org.opensearch.index.compositeindex.startree.aggregators; import org.opensearch.index.compositeindex.MetricType; +import org.opensearch.index.compositeindex.startree.data.DataType; import org.opensearch.test.OpenSearchTestCase; public class ValueAggregatorFactoryTests extends OpenSearchTestCase {