From bc6fe9b13156d9978cd1133c602ec196699d1516 Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Wed, 25 Sep 2024 18:48:04 +0530 Subject: [PATCH 1/2] star tree keyword changes Signed-off-by: Bharathwaj G --- .../index/mapper/StarTreeMapperIT.java | 32 +++++--- .../Composite99DocValuesWriter.java | 54 +++++++++---- .../datacube/DimensionFactory.java | 28 ++++--- .../datacube/DimensionType.java | 8 +- .../datacube/KeywordDimension.java | 75 +++++++++++++++++++ .../builder/StarTreeDocsFileManager.java | 3 + .../index/mapper/KeywordFieldMapper.java | 7 ++ .../index/mapper/StarTreeMapperTests.java | 43 ++++++++++- 8 files changed, 215 insertions(+), 35 deletions(-) create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java diff --git a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java index 5840884f5422a..c91c4d7bbb63b 100644 --- a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java @@ -56,7 +56,7 @@ public class StarTreeMapperIT extends OpenSearchIntegTestCase { .put(IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING.getKey(), new ByteSizeValue(512, ByteSizeUnit.MB)) .build(); - private static XContentBuilder createMinimalTestMapping(boolean invalidDim, boolean invalidMetric, boolean keywordDim) { + private static XContentBuilder createMinimalTestMapping(boolean invalidDim, boolean invalidMetric, boolean ipdim) { try { return jsonBuilder().startObject() .startObject("composite") @@ -68,12 +68,15 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool .endObject() .startArray("ordered_dimensions") .startObject() - .field("name", getDim(invalidDim, keywordDim)) + .field("name", getDim(invalidDim, ipdim)) + .endObject() + .startObject() + .field("name", "keyword_dv") .endObject() .endArray() .startArray("metrics") .startObject() - .field("name", getDim(invalidMetric, false)) + .field("name", getMetric(invalidMetric, false)) .endObject() .endArray() .endObject() @@ -99,6 +102,10 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool .field("type", "keyword") .field("doc_values", false) .endObject() + .startObject("ip") + .field("type", "ip") + .field("doc_values", false) + .endObject() .endObject() .endObject(); } catch (IOException e) { @@ -356,10 +363,19 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo } private static String getDim(boolean hasDocValues, boolean isKeyword) { + if (hasDocValues) { + return random().nextBoolean() ? "numeric" : "keyword"; + } else if (isKeyword) { + return "ip"; + } + return "numeric_dv"; + } + + private static String getMetric(boolean hasDocValues, boolean isKeyword) { if (hasDocValues) { return "numeric"; } else if (isKeyword) { - return "keyword"; + return "ip"; } return "numeric_dv"; } @@ -398,6 +414,7 @@ public void testValidCompositeIndex() { assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName()); } assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("keyword_dv", starTreeFieldType.getDimensions().get(2).getField()); assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField()); List expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG); assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics()); @@ -665,10 +682,7 @@ public void testInvalidDimCompositeIndex() { IllegalArgumentException.class, () -> prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createMinimalTestMapping(true, false, false)).get() ); - assertEquals( - "Aggregations not supported for the dimension field [numeric] with field type [integer] as part of star tree field", - ex.getMessage() - ); + assertTrue(ex.getMessage().startsWith("Aggregations not supported for the dimension field ")); } public void testMaxDimsCompositeIndex() { @@ -734,7 +748,7 @@ public void testUnsupportedDim() { () -> prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createMinimalTestMapping(false, false, true)).get() ); assertEquals( - "Failed to parse mapping [_doc]: unsupported field type associated with dimension [keyword] as part of star tree field [startree-1]", + "Failed to parse mapping [_doc]: unsupported field type associated with dimension [ip] as part of star tree field [startree-1]", ex.getMessage() ); } diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java index 0d4e35f7c3ab8..3c914a2b7ad0a 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java @@ -21,6 +21,7 @@ import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; import org.opensearch.common.annotation.ExperimentalApi; @@ -34,6 +35,7 @@ import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.mapper.CompositeMappedFieldType; import org.opensearch.index.mapper.DocCountFieldMapper; +import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.StarTreeMapper; @@ -82,14 +84,7 @@ public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState this.compositeMappedFieldTypes = mapperService.getCompositeFieldTypes(); compositeFieldSet = new HashSet<>(); segmentFieldSet = new HashSet<>(); - // TODO : add integ test for this - for (FieldInfo fi : this.state.fieldInfos) { - if (DocValuesType.SORTED_NUMERIC.equals(fi.getDocValuesType())) { - segmentFieldSet.add(fi.name); - } else if (fi.name.equals(DocCountFieldMapper.NAME)) { - segmentFieldSet.add(fi.name); - } - } + addStarTreeSupportedFieldsFromSegment(); for (CompositeMappedFieldType type : compositeMappedFieldTypes) { compositeFieldSet.addAll(type.fields()); } @@ -148,6 +143,19 @@ public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState segmentHasCompositeFields = Collections.disjoint(segmentFieldSet, compositeFieldSet) == false; } + private void addStarTreeSupportedFieldsFromSegment() { + // TODO : add integ test for this + for (FieldInfo fi : this.state.fieldInfos) { + if (DocValuesType.SORTED_NUMERIC.equals(fi.getDocValuesType())) { + segmentFieldSet.add(fi.name); + } else if (DocValuesType.SORTED_SET.equals(fi.getDocValuesType())) { + segmentFieldSet.add(fi.name); + } else if (fi.name.equals(DocCountFieldMapper.NAME)) { + segmentFieldSet.add(fi.name); + } + } + } + @Override public void addNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { delegate.addNumericField(field, valuesProducer); @@ -179,6 +187,10 @@ public void addSortedNumericField(FieldInfo field, DocValuesProducer valuesProdu @Override public void addSortedSetField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { delegate.addSortedSetField(field, valuesProducer); + // Perform this only during flush flow + if (mergeState.get() == null && segmentHasCompositeFields) { + createCompositeIndicesIfPossible(valuesProducer, field); + } } @Override @@ -235,6 +247,7 @@ private void createCompositeIndicesIfPossible(DocValuesProducer valuesProducer, * Add empty doc values for fields not present in segment */ private void addDocValuesForEmptyField(String compositeField) { + // special case for doc count if (compositeField.equals(DocCountFieldMapper.NAME)) { fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() { @Override @@ -243,16 +256,29 @@ public NumericDocValues getNumeric(FieldInfo field) { } }); } else { - fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() { - @Override - public SortedNumericDocValues getSortedNumeric(FieldInfo field) { - return DocValues.emptySortedNumeric(); - } - }); + if (isSortedSetField(compositeField)) { + fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() { + @Override + public SortedSetDocValues getSortedSet(FieldInfo field) { + return DocValues.emptySortedSet(); + } + }); + } else { + fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() { + @Override + public SortedNumericDocValues getSortedNumeric(FieldInfo field) { + return DocValues.emptySortedNumeric(); + } + }); + } } compositeFieldSet.remove(compositeField); } + private boolean isSortedSetField(String field) { + return mapperService.fieldType(field) instanceof KeywordFieldMapper.KeywordFieldType; + } + @Override public void merge(MergeState mergeState) throws IOException { this.mergeState.compareAndSet(null, mergeState); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java index 7e72a3f0d9de6..e834706e2fa9d 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java @@ -24,6 +24,7 @@ import java.util.stream.Collectors; import static org.opensearch.index.compositeindex.datacube.DateDimension.CALENDAR_INTERVALS; +import static org.opensearch.index.compositeindex.datacube.KeywordDimension.KEYWORD; /** * Dimension factory class mainly used to parse and create dimension from the mappings @@ -43,6 +44,8 @@ public static Dimension parseAndCreateDimension( return parseAndCreateDateDimension(name, dimensionMap, c); case NumericDimension.NUMERIC: return new NumericDimension(name); + case KEYWORD: + return new KeywordDimension(name); default: throw new IllegalArgumentException( String.format(Locale.ROOT, "unsupported field type associated with dimension [%s] as part of star tree field", name) @@ -56,16 +59,23 @@ public static Dimension parseAndCreateDimension( Map dimensionMap, Mapper.TypeParser.ParserContext c ) { - if (builder.getSupportedDataCubeDimensionType().isPresent() - && builder.getSupportedDataCubeDimensionType().get().equals(DimensionType.DATE)) { - return parseAndCreateDateDimension(name, dimensionMap, c); - } else if (builder.getSupportedDataCubeDimensionType().isPresent() - && builder.getSupportedDataCubeDimensionType().get().equals(DimensionType.NUMERIC)) { + if (builder.getSupportedDataCubeDimensionType().isEmpty()) { + throw new IllegalArgumentException( + String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name) + ); + } + switch (builder.getSupportedDataCubeDimensionType().get()) { + case DATE: + return parseAndCreateDateDimension(name, dimensionMap, c); + case NUMERIC: return new NumericDimension(name); - } - throw new IllegalArgumentException( - String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name) - ); + case KEYWORD: + return new KeywordDimension(name); + default: + throw new IllegalArgumentException( + String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name) + ); + } } private static DateDimension parseAndCreateDateDimension( diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java index 4b9faea331752..d327f8ca1fa1e 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java @@ -27,5 +27,11 @@ public enum DimensionType { * Represents a date dimension type. * This is used for dimensions that contain date or timestamp values. */ - DATE + DATE, + + /** + * Represents a keyword dimension type. + * This is used for dimensions that contain keyword ordinals. + */ + KEYWORD } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java new file mode 100644 index 0000000000000..dd6cedade43ff --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java @@ -0,0 +1,75 @@ +/* + * 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; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.index.mapper.CompositeDataCubeFieldType; + +import java.io.IOException; +import java.util.List; +import java.util.Objects; + +/** + * Composite index keyword dimension class + * + * @opensearch.experimental + */ +@ExperimentalApi +public class KeywordDimension implements Dimension { + public static final String KEYWORD = "keyword"; + private final String field; + + public KeywordDimension(String field) { + this.field = field; + } + + @Override + public String getField() { + return field; + } + + @Override + public int getNumSubDimensions() { + return 1; + } + + @Override + public int setDimensionValues(Long value, Long[] dims, int index) { + dims[index++] = value; + return index; + } + + @Override + public List getDimensionFieldsNames() { + return List.of(field); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field(CompositeDataCubeFieldType.NAME, field); + builder.field(CompositeDataCubeFieldType.TYPE, KEYWORD); + builder.endObject(); + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + KeywordDimension dimension = (KeywordDimension) o; + return Objects.equals(field, dimension.getField()); + } + + @Override + public int hashCode() { + return Objects.hash(field); + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java index 7e920b912731d..98c3e5c6d71e6 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java @@ -14,6 +14,7 @@ import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.store.RandomAccessInput; +import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; @@ -45,7 +46,9 @@ *

The set of 'star-tree.documents' files is maintained, and a tracker array is used to keep track of the start document ID for each file. * Once the number of files reaches a set threshold, the files are merged. * + * @opensearch.experimental */ +@ExperimentalApi public class StarTreeDocsFileManager extends AbstractDocumentsFileManager implements Closeable { private static final Logger logger = LogManager.getLogger(StarTreeDocsFileManager.class); private static final String STAR_TREE_DOC_FILE_NAME = "star-tree.documents"; diff --git a/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java index 11ff601b3fd6d..13577d3c3aa56 100644 --- a/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java @@ -59,6 +59,7 @@ import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.analysis.IndexAnalyzers; import org.opensearch.index.analysis.NamedAnalyzer; +import org.opensearch.index.compositeindex.datacube.DimensionType; import org.opensearch.index.fielddata.IndexFieldData; import org.opensearch.index.fielddata.plain.SortedSetOrdinalsIndexFieldData; import org.opensearch.index.query.QueryShardContext; @@ -73,6 +74,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.function.Supplier; import static org.opensearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES; @@ -254,6 +256,11 @@ public KeywordFieldMapper build(BuilderContext context) { this ); } + + @Override + public Optional getSupportedDataCubeDimensionType() { + return Optional.of(DimensionType.KEYWORD); + } } public static final TypeParser PARSER = new TypeParser((n, c) -> new Builder(n, c.getIndexAnalyzers())); diff --git a/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java index aac460bd5e332..8ec34b3eb660c 100644 --- a/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java @@ -672,6 +672,9 @@ private XContentBuilder getExpandedMappingWithJustAvg(String dim, String metric) b.startObject("size"); b.field("type", "integer"); b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -718,6 +721,7 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo .field("type", "integer") .field("doc_values", true) .endObject() + .endObject() .endObject(); } catch (IOException e) { @@ -772,6 +776,9 @@ private XContentBuilder getExpandedMappingWithJustSum(String dim, String metric) b.startObject("size"); b.field("type", "integer"); b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -823,6 +830,9 @@ private XContentBuilder getExpandedMappingWithSumAndCount(String dim, String met b.startObject("size"); b.field("type", "integer"); b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -866,6 +876,9 @@ private XContentBuilder getMinMappingWithDateDims(boolean calendarIntervalsExcee b.startObject(); b.field("name", "metric_field"); b.endObject(); + b.startObject(); + b.field("name", "keyword1"); + b.endObject(); } b.endArray(); @@ -895,6 +908,9 @@ private XContentBuilder getMinMappingWithDateDims(boolean calendarIntervalsExcee b.startObject("metric_field"); b.field("type", "integer"); b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); @@ -920,6 +936,9 @@ private XContentBuilder getMinMapping( b.startObject(); b.field("name", "status"); b.endObject(); + b.startObject(); + b.field("name", "keyword1"); + b.endObject(); b.endArray(); } if (!isEmptyMetrics) { @@ -951,6 +970,9 @@ private XContentBuilder getMinMapping( b.field("type", "integer"); b.endObject(); } + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -1018,7 +1040,9 @@ private XContentBuilder getMinMappingWith2StarTrees() throws IOException { b.startObject("metric_field"); b.field("type", "integer"); b.endObject(); - + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -1058,6 +1082,9 @@ private XContentBuilder getInvalidMapping( b.startObject(); b.field("name", "status"); b.endObject(); + b.startObject(); + b.field("name", "keyword1"); + b.endObject(); } b.endArray(); b.startArray("metrics"); @@ -1090,7 +1117,7 @@ private XContentBuilder getInvalidMapping( if (!invalidDimType) { b.field("type", "integer"); } else { - b.field("type", "keyword"); + b.field("type", "ip"); } b.endObject(); b.startObject("metric_field"); @@ -1100,6 +1127,9 @@ private XContentBuilder getInvalidMapping( b.field("type", "integer"); } b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -1132,6 +1162,9 @@ private XContentBuilder getInvalidMappingWithDv( b.startObject(); b.field("name", "status"); b.endObject(); + b.startObject(); + b.field("name", "keyword1"); + b.endObject(); } b.endArray(); b.startArray("metrics"); @@ -1168,6 +1201,9 @@ private XContentBuilder getInvalidMappingWithDv( b.field("doc_values", "true"); } b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }); } @@ -1224,6 +1260,9 @@ public void testEmptyName() { b.startObject("status"); b.field("type", "integer"); b.endObject(); + b.startObject("keyword1"); + b.field("type", "keyword"); + b.endObject(); b.endObject(); }))); assertThat(e.getMessage(), containsString("name cannot be empty string")); From 57b0517a3ebc66724dd24f098d7cd67c96993552 Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Fri, 4 Oct 2024 10:36:30 +0530 Subject: [PATCH 2/2] Complete keyword changes Signed-off-by: Bharathwaj G --- .../lucene/index/DocValuesWriterWrapper.java | 18 ++ .../SortedNumericDocValuesWriterWrapper.java | 3 +- .../SortedSetDocValuesWriterWrapper.java | 58 +++++ .../opensearch/common/util/FeatureFlags.java | 2 +- .../Composite99DocValuesReader.java | 22 +- .../Composite99DocValuesWriter.java | 15 +- .../datacube/DateDimension.java | 6 + .../compositeindex/datacube/Dimension.java | 3 + .../datacube/KeywordDimension.java | 14 +- .../datacube/NumericDimension.java | 6 + .../datacube/ReadDimension.java | 14 + .../datacube/startree/StarTreeField.java | 12 +- .../startree/builder/BaseStarTreeBuilder.java | 246 +++++++++++++++--- .../builder/OffHeapStarTreeBuilder.java | 23 +- .../builder/OnHeapStarTreeBuilder.java | 30 ++- .../startree/builder/StarTreeBuilder.java | 6 +- .../startree/builder/StarTreesBuilder.java | 25 +- .../fileformats/meta/StarTreeMetadata.java | 46 +++- .../meta/StarTreeMetadataWriter.java | 26 +- .../startree/index/StarTreeValues.java | 37 ++- .../utils/SequentialDocValuesIterator.java | 29 +++ .../startree/utils/StarTreeUtils.java | 12 +- .../SortedSetStarTreeValuesIterator.java | 59 +++++ .../StarTreeDocValuesFormatTests.java | 24 +- .../builder/BaseStarTreeBuilderTests.java | 8 +- .../startree/builder/BuilderTestsUtils.java | 4 +- .../builder/StarTreeBuildMetricTests.java | 21 +- .../StarTreeBuilderFlushFlowTests.java | 18 +- .../StarTreeBuilderMergeFlowTests.java | 87 ++++--- .../builder/StarTreeBuilderTestCase.java | 22 +- .../meta/StarTreeMetadataTests.java | 8 +- .../startree/utils/StarTreeUtilsTests.java | 3 +- 32 files changed, 743 insertions(+), 164 deletions(-) create mode 100644 server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java create mode 100644 server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java diff --git a/server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java new file mode 100644 index 0000000000000..3b8354c074bcb --- /dev/null +++ b/server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java @@ -0,0 +1,18 @@ +/* + * 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.search.DocIdSetIterator; + +/** + * Base wrapper class for DocValuesWriter. + */ +public abstract class DocValuesWriterWrapper { + public abstract T getDocValues(); +} diff --git a/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java index f7759fcced284..36fd29c7fd580 100644 --- a/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java +++ b/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java @@ -18,7 +18,7 @@ * * @opensearch.experimental */ -public class SortedNumericDocValuesWriterWrapper { +public class SortedNumericDocValuesWriterWrapper extends DocValuesWriterWrapper { private final SortedNumericDocValuesWriter sortedNumericDocValuesWriter; @@ -47,6 +47,7 @@ public void addValue(int docID, long value) { * * @return the {@link SortedNumericDocValues} instance */ + @Override public SortedNumericDocValues getDocValues() { return sortedNumericDocValuesWriter.getDocValues(); } diff --git a/server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java new file mode 100644 index 0000000000000..229617b4a44b0 --- /dev/null +++ b/server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java @@ -0,0 +1,58 @@ +/* + * 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.util.ByteBlockPool; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.Counter; + +/** + * A wrapper class for writing sorted set doc values. + *

+ * This class provides a convenient way to add sorted set doc values to a field + * and retrieve the corresponding {@link SortedSetDocValues} instance. + * + * @opensearch.experimental + */ +public class SortedSetDocValuesWriterWrapper extends DocValuesWriterWrapper { + + private final SortedSetDocValuesWriter sortedSetDocValuesWriterWrapper; + + /** + * Sole constructor. Constructs a new {@link SortedSetDocValuesWriterWrapper} instance. + * + * @param fieldInfo the field information for the field being written + * @param counter a counter for tracking memory usage + * @param byteBlockPool a byte block pool for allocating byte blocks + * @see SortedSetDocValuesWriter + */ + public SortedSetDocValuesWriterWrapper(FieldInfo fieldInfo, Counter counter, ByteBlockPool byteBlockPool) { + sortedSetDocValuesWriterWrapper = new SortedSetDocValuesWriter(fieldInfo, counter, byteBlockPool); + } + + /** + * Adds a bytes ref value to the sorted set doc values for the specified document. + * + * @param docID the document ID + * @param value the value to add + */ + public void addValue(int docID, BytesRef value) { + sortedSetDocValuesWriterWrapper.addValue(docID, value); + } + + /** + * Returns the {@link SortedSetDocValues} instance containing the sorted numeric doc values + * + * @return the {@link SortedSetDocValues} instance + */ + @Override + public SortedSetDocValues getDocValues() { + return sortedSetDocValuesWriterWrapper.getDocValues(); + } +} diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index 6df68013a8119..e663d8429da13 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -100,7 +100,7 @@ public class FeatureFlags { * aggregations. */ public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled"; - public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope); + public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope); /** * Gates the functionality of application based configuration templates. diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java index 7901336151c8e..8966ab5db8317 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java @@ -15,6 +15,7 @@ import org.apache.lucene.index.BinaryDocValues; import org.apache.lucene.index.CorruptIndexException; import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.IndexFileNames; @@ -40,6 +41,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -111,7 +113,7 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r readState.segmentInfo.getId(), readState.segmentSuffix ); - + Map dimensionFieldTypeMap = new HashMap<>(); while (true) { // validate magic marker @@ -155,13 +157,15 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r compositeIndexInputMap.put(compositeFieldName, starTreeIndexInput); compositeIndexMetadataMap.put(compositeFieldName, starTreeMetadata); - List dimensionFields = starTreeMetadata.getDimensionFields(); - + Map dimensionFieldToDocValuesMap = starTreeMetadata.getDimensionFields(); // generating star tree unique fields (fully qualified name for dimension and metrics) - for (String dimensions : dimensionFields) { - fields.add(fullyQualifiedFieldNameForStarTreeDimensionsDocValues(compositeFieldName, dimensions)); + for (Map.Entry dimensionEntry : dimensionFieldToDocValuesMap.entrySet()) { + String dimName = fullyQualifiedFieldNameForStarTreeDimensionsDocValues( + compositeFieldName, + dimensionEntry.getKey() + ); + fields.add(dimName); } - // adding metric fields for (Metric metric : starTreeMetadata.getMetrics()) { for (MetricStat metricStat : metric.getBaseMetrics()) { @@ -184,7 +188,7 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r // populates the dummy list of field infos to fetch doc id set iterators for respective fields. // the dummy field info is used to fetch the doc id set iterators for respective fields based on field name - FieldInfos fieldInfos = new FieldInfos(getFieldInfoList(fields)); + FieldInfos fieldInfos = new FieldInfos(getFieldInfoList(fields, dimensionFieldTypeMap)); this.readState = new SegmentReadState(readState.directory, readState.segmentInfo, fieldInfos, readState.context); // initialize star-tree doc values producer @@ -298,4 +302,8 @@ public static SortedNumericDocValues getSortedNumericDocValues(SortedNumericDocV return sortedNumeric == null ? DocValues.emptySortedNumeric() : sortedNumeric; } + public static SortedSetDocValues getSortedSetDocValues(SortedSetDocValues sortedSetDv) { + return sortedSetDv == null ? DocValues.emptySortedSet() : sortedSetDv; + } + } diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java index 3c914a2b7ad0a..6fc2d6f463642 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java @@ -73,6 +73,7 @@ public class Composite99DocValuesWriter extends DocValuesConsumer { private final AtomicInteger fieldNumberAcrossCompositeFields; private final Map fieldProducerMap = new HashMap<>(); + private final Map fieldDocIdSetIteratorMap = new HashMap<>(); public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState segmentWriteState, MapperService mapperService) throws IOException { @@ -191,6 +192,11 @@ public void addSortedSetField(FieldInfo field, DocValuesProducer valuesProducer) if (mergeState.get() == null && segmentHasCompositeFields) { createCompositeIndicesIfPossible(valuesProducer, field); } + if (mergeState.get() != null) { + if (compositeFieldSet.contains(field.name)) { + fieldDocIdSetIteratorMap.put(field.name, valuesProducer.getSortedSet(field)); + } + } } @Override @@ -339,7 +345,14 @@ private void mergeStarTreeFields(MergeState mergeState) throws IOException { } } try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService, fieldNumberAcrossCompositeFields)) { - starTreesBuilder.buildDuringMerge(metaOut, dataOut, starTreeSubsPerField, composite99DocValuesConsumer); + starTreesBuilder.buildDuringMerge( + metaOut, + dataOut, + starTreeSubsPerField, + composite99DocValuesConsumer, + mergeState, + fieldDocIdSetIteratorMap + ); } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java index ee6d5b4680c73..88a67e1134067 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube; +import org.apache.lucene.index.DocValuesType; import org.opensearch.common.Rounding; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.time.DateUtils; @@ -98,6 +99,11 @@ public List getSubDimensionNames() { return fields; } + @Override + public DocValuesType getDocValuesType() { + return DocValuesType.SORTED_NUMERIC; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject("date_dimension"); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java index cfa8d3a2a8164..3d71b38881693 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube; +import org.apache.lucene.index.DocValuesType; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.ToXContent; @@ -42,4 +43,6 @@ public interface Dimension extends ToXContent { * Returns the list of dimension fields that represent the dimension */ List getSubDimensionNames(); + + DocValuesType getDocValuesType(); } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java index dd6cedade43ff..60135244e2cc4 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube; +import org.apache.lucene.index.DocValuesType; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.mapper.CompositeDataCubeFieldType; @@ -15,6 +16,7 @@ import java.io.IOException; import java.util.List; import java.util.Objects; +import java.util.function.Consumer; /** * Composite index keyword dimension class @@ -41,16 +43,20 @@ public int getNumSubDimensions() { } @Override - public int setDimensionValues(Long value, Long[] dims, int index) { - dims[index++] = value; - return index; + public void setDimensionValues(Long value, Consumer dimSetter) { + dimSetter.accept(value); } @Override - public List getDimensionFieldsNames() { + public List getSubDimensionNames() { return List.of(field); } + @Override + public DocValuesType getDocValuesType() { + return DocValuesType.SORTED_SET; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java index acc14f5f05c68..fe9e3d17c0047 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube; +import org.apache.lucene.index.DocValuesType; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.mapper.CompositeDataCubeFieldType; @@ -50,6 +51,11 @@ public List getSubDimensionNames() { return List.of(field); } + @Override + public DocValuesType getDocValuesType() { + return DocValuesType.SORTED_NUMERIC; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java index be3667f10b6da..220c1b97cb4df 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube; +import org.apache.lucene.index.DocValuesType; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.mapper.CompositeDataCubeFieldType; @@ -24,9 +25,16 @@ public class ReadDimension implements Dimension { public static final String READ = "read"; private final String field; + private final DocValuesType docValuesType; public ReadDimension(String field) { this.field = field; + this.docValuesType = DocValuesType.SORTED_NUMERIC; + } + + public ReadDimension(String field, DocValuesType docValuesType) { + this.field = field; + this.docValuesType = docValuesType; } public String getField() { @@ -48,6 +56,11 @@ public List getSubDimensionNames() { return List.of(field); } + @Override + public DocValuesType getDocValuesType() { + return docValuesType; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); @@ -69,4 +82,5 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(field); } + } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java index 833bf63c04a18..37b59fc1f59c8 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java @@ -8,6 +8,7 @@ package org.opensearch.index.compositeindex.datacube.startree; +import org.apache.lucene.index.DocValuesType; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; @@ -33,6 +34,7 @@ public class StarTreeField implements ToXContent { private final List metrics; private final StarTreeFieldConfiguration starTreeConfig; private final List dimensionNames; + private final List dimensionDocValueTypes; private final List metricNames; public StarTreeField(String name, List dimensions, List metrics, StarTreeFieldConfiguration starTreeConfig) { @@ -41,8 +43,12 @@ public StarTreeField(String name, List dimensions, List metri this.metrics = metrics; this.starTreeConfig = starTreeConfig; dimensionNames = new ArrayList<>(); + dimensionDocValueTypes = new ArrayList<>(); for (Dimension dimension : dimensions) { - dimensionNames.addAll(dimension.getSubDimensionNames()); + for (String dimensionName : dimension.getSubDimensionNames()) { + dimensionNames.add(dimensionName); + dimensionDocValueTypes.add(dimension.getDocValuesType()); + } } metricNames = new ArrayList<>(); for (Metric metric : metrics) { @@ -64,6 +70,10 @@ public List getDimensionNames() { return dimensionNames; } + public List getDimensionDocValueTypes() { + return dimensionDocValueTypes; + } + public List getMetricNames() { return metricNames; } 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 2d4938eeb45b3..b61016414e2ee 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 @@ -13,15 +13,28 @@ import org.apache.lucene.codecs.DocValuesProducer; import org.apache.lucene.index.DocValues; import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.DocValuesWriterWrapper; import org.apache.lucene.index.EmptyDocValuesProducer; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FilteredTermsEnum; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.OrdinalMap; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.index.SortedNumericDocValuesWriterWrapper; +import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.index.SortedSetDocValuesWriterWrapper; +import org.apache.lucene.index.TermsEnum; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.ByteBlockPool; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.Counter; +import org.apache.lucene.util.LongBitSet; +import org.apache.lucene.util.LongValues; import org.apache.lucene.util.NumericUtils; +import org.apache.lucene.util.packed.PackedInts; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.Metric; import org.opensearch.index.compositeindex.datacube.MetricStat; @@ -36,6 +49,7 @@ 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.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator; import org.opensearch.index.mapper.DocCountFieldMapper; import org.opensearch.index.mapper.FieldMapper; import org.opensearch.index.mapper.FieldValueConverter; @@ -45,6 +59,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; @@ -58,6 +73,7 @@ import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeDimensionsDocValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.getFieldInfo; +import static org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator.NO_MORE_ENTRIES; import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.DOUBLE; import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.LONG; @@ -86,12 +102,13 @@ public abstract class BaseStarTreeBuilder implements StarTreeBuilder { protected final int maxLeafDocuments; List dimensionsSplitOrder = new ArrayList<>(); protected final InMemoryTreeNode rootNode = getNewNode(); - protected final StarTreeField starTreeField; private final SegmentWriteState writeState; private final IndexOutput metaOut; private final IndexOutput dataOut; + private final Counter bytesUsed = Counter.newCounter(); + Map sortedSetDocValuesMap = new HashMap<>(); /** * Reads all the configuration related to dimensions and metrics, builds a star-tree based on the different construction parameters. @@ -236,11 +253,23 @@ public void build( String dimension = dimensionsSplitOrder.get(i).getField(); FieldInfo dimensionFieldInfo = writeState.fieldInfos.fieldInfo(dimension); if (dimensionFieldInfo == null) { - dimensionFieldInfo = getFieldInfo(dimension, DocValuesType.SORTED_NUMERIC); + // TODO : make this better + dimensionFieldInfo = getFieldInfo(dimension, dimensionsSplitOrder.get(i).getDocValuesType()); } - dimensionReaders[i] = new SequentialDocValuesIterator( - new SortedNumericStarTreeValuesIterator(fieldProducerMap.get(dimensionFieldInfo.name).getSortedNumeric(dimensionFieldInfo)) + dimensionReaders[i] = getSequentialDocValuesIterator( + dimensionFieldInfo, + fieldProducerMap, + dimensionsSplitOrder.get(i).getDocValuesType() ); + + if (dimensionsSplitOrder.get(i).getDocValuesType().equals(DocValuesType.SORTED_SET)) { + // This is needed as we need to write the ordinals and also the bytesRef associated with it + sortedSetDocValuesMap.put( + dimensionsSplitOrder.get(i).getField(), + fieldProducerMap.get(dimensionFieldInfo.name).getSortedSet(dimensionFieldInfo) + ); + } + } Iterator starTreeDocumentIterator = sortAndAggregateSegmentDocuments(dimensionReaders, metricReaders); logger.debug("Sorting and aggregating star-tree in ms : {}", (System.currentTimeMillis() - startTime)); @@ -248,6 +277,95 @@ public void build( logger.debug("Finished Building star-tree in ms : {}", (System.currentTimeMillis() - startTime)); } + /** + * Copy of BitsFilteredTermsEnum from DocValuesConsumer + */ + static class BitsFilteredTermsEnum extends FilteredTermsEnum { + final LongBitSet liveTerms; + + BitsFilteredTermsEnum(TermsEnum in, LongBitSet liveTerms) { + super(in, false); + assert liveTerms != null; + this.liveTerms = liveTerms; + } + + @Override + protected AcceptStatus accept(BytesRef term) throws IOException { + if (liveTerms.get(ord())) { + return AcceptStatus.YES; + } else { + return AcceptStatus.NO; + } + } + } + + /** + * Returns the sequential doc values iterator for the given field based on associated docValuesType + */ + private SequentialDocValuesIterator getSequentialDocValuesIterator( + FieldInfo fieldInfo, + Map fieldProducerMap, + DocValuesType type + ) throws IOException { + switch (type) { + case SORTED_NUMERIC: + return new SequentialDocValuesIterator( + new SortedNumericStarTreeValuesIterator(fieldProducerMap.get(fieldInfo.name).getSortedNumeric(fieldInfo)) + ); + case SORTED_SET: + return new SequentialDocValuesIterator( + new SortedSetStarTreeValuesIterator(fieldProducerMap.get(fieldInfo.name).getSortedSet(fieldInfo)) + ); + default: + throw new IllegalArgumentException("Unsupported type: " + type); + } + } + + /** + * Returns the ordinal map based on given star-tree values across different segments + */ + protected OrdinalMap getOrdinalMap(List starTreeValuesSubs, MergeState mergeState) throws IOException { + long curr = System.currentTimeMillis(); + List toMerge = new ArrayList<>(); + + for (StarTreeValues starTree : starTreeValuesSubs) { + for (String dimName : starTree.getStarTreeField().getDimensionNames()) { + if (starTree.getDimensionValuesIterator(dimName) instanceof SortedSetStarTreeValuesIterator) { + toMerge.add((SortedSetStarTreeValuesIterator) starTree.getDimensionValuesIterator(dimName)); + } + } + } + if (toMerge.isEmpty()) return null; + // step 1: iterate through each sub and mark terms still in use + TermsEnum[] liveTerms = new TermsEnum[toMerge.size()]; + long[] weights = new long[liveTerms.length]; + for (int sub = 0; sub < liveTerms.length; sub++) { + SortedSetStarTreeValuesIterator dv = toMerge.get(sub); + Bits liveDocs = mergeState.liveDocs[sub]; + if (liveDocs == null) { + liveTerms[sub] = dv.termsEnum(); + weights[sub] = dv.getValueCount(); + } else { + LongBitSet bitset = new LongBitSet(dv.getValueCount()); + int docID; + while ((docID = dv.nextEntry()) != NO_MORE_ENTRIES) { + if (liveDocs.get(docID)) { + for (int i = 0; i < dv.docValueCount(); i++) { + bitset.set(dv.nextOrd()); + } + } + } + liveTerms[sub] = new BitsFilteredTermsEnum(dv.termsEnum(), bitset); + weights[sub] = bitset.cardinality(); + } + } + // step 2: create ordinal map (this conceptually does the "merging") + final OrdinalMap map = OrdinalMap.build(null, liveTerms, weights, PackedInts.COMPACT); + logger.debug("Map size in bytes : {}", map.ramBytesUsed()); + logger.debug("Ordinal map takes : {} ", System.currentTimeMillis() - curr); + return map; + } + /** * Builds the star tree using sorted and aggregated star-tree Documents * @@ -265,7 +383,7 @@ public void build( appendDocumentsToStarTree(starTreeDocumentIterator); int numStarTreeDocument = numStarTreeDocs; - logger.debug("Generated star tree docs : [{}] from segment docs : [{}]", numStarTreeDocument, numSegmentStarTreeDocument); + logger.info("Generated star tree docs : [{}] from segment docs : [{}]", numStarTreeDocument, numSegmentStarTreeDocument); if (numStarTreeDocs == 0) { // serialize the star tree data @@ -275,7 +393,7 @@ public void build( constructStarTree(rootNode, 0, numStarTreeDocs); int numStarTreeDocumentUnderStarNode = numStarTreeDocs - numStarTreeDocument; - logger.debug( + logger.info( "Finished constructing star-tree, got [ {} ] tree nodes and [ {} ] starTreeDocument under star-node", numStarTreeNodes, numStarTreeDocumentUnderStarNode @@ -283,7 +401,7 @@ public void build( createAggregatedDocs(rootNode); int numAggregatedStarTreeDocument = numStarTreeDocs - numStarTreeDocument - numStarTreeDocumentUnderStarNode; - logger.debug("Finished creating aggregated documents : {}", numAggregatedStarTreeDocument); + logger.info("Finished creating aggregated documents : {}", numAggregatedStarTreeDocument); // Create doc values indices in disk createSortedDocValuesIndices(starTreeDocValuesConsumer, fieldNumberAcrossStarTrees); @@ -298,6 +416,9 @@ void appendDocumentsToStarTree(Iterator starTreeDocumentIterat } } + /** + * Writes star tree structure to file format + */ private void serializeStarTree(int numSegmentStarTreeDocuments, int numStarTreeDocs) throws IOException { // serialize the star tree data long dataFilePointer = dataOut.getFilePointer(); @@ -317,10 +438,13 @@ private void serializeStarTree(int numSegmentStarTreeDocuments, int numStarTreeD ); } + /** + * Creates the star-tree docValues indices in disk + */ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, AtomicInteger fieldNumberAcrossStarTrees) throws IOException { - List dimensionWriters = new ArrayList<>(); - List metricWriters = new ArrayList<>(); + List> dimensionWriters = new ArrayList<>(); + List> metricWriters = new ArrayList<>(); FieldInfo[] dimensionFieldInfoList = new FieldInfo[numDimensions]; FieldInfo[] metricFieldInfoList = new FieldInfo[metricAggregatorInfos.size()]; int dimIndex = 0; @@ -328,16 +452,21 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A for (String name : dim.getSubDimensionNames()) { final FieldInfo fi = getFieldInfo( fullyQualifiedFieldNameForStarTreeDimensionsDocValues(starTreeField.getName(), name), - DocValuesType.SORTED_NUMERIC, + dim.getDocValuesType(), fieldNumberAcrossStarTrees.getAndIncrement() ); dimensionFieldInfoList[dimIndex] = fi; - dimensionWriters.add(new SortedNumericDocValuesWriterWrapper(fi, Counter.newCounter())); + if (dim.getDocValuesType().equals(DocValuesType.SORTED_SET)) { + ByteBlockPool.DirectTrackingAllocator byteBlockAllocator = new ByteBlockPool.DirectTrackingAllocator(bytesUsed); + ByteBlockPool docValuesBytePool = new ByteBlockPool(byteBlockAllocator); + dimensionWriters.add(new SortedSetDocValuesWriterWrapper(fi, bytesUsed, docValuesBytePool)); + } else { + dimensionWriters.add(new SortedNumericDocValuesWriterWrapper(fi, bytesUsed)); + } dimIndex++; } } for (int i = 0; i < metricAggregatorInfos.size(); i++) { - final FieldInfo fi = getFieldInfo( fullyQualifiedFieldNameForStarTreeMetricsDocValues( starTreeField.getName(), @@ -347,16 +476,19 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A DocValuesType.SORTED_NUMERIC, fieldNumberAcrossStarTrees.getAndIncrement() ); - metricFieldInfoList[i] = fi; - metricWriters.add(new SortedNumericDocValuesWriterWrapper(fi, Counter.newCounter())); + metricWriters.add(new SortedNumericDocValuesWriterWrapper(fi, bytesUsed)); } for (int docId = 0; docId < numStarTreeDocs; docId++) { StarTreeDocument starTreeDocument = getStarTreeDocument(docId); - for (int i = 0; i < starTreeDocument.dimensions.length; i++) { - if (starTreeDocument.dimensions[i] != null) { - dimensionWriters.get(i).addValue(docId, starTreeDocument.dimensions[i]); + int idx = 0; + for (Dimension dim : dimensionsSplitOrder) { + for (String name : dim.getSubDimensionNames()) { + if (starTreeDocument.dimensions[idx] != null) { + indexDocValue(dimensionWriters.get(idx), docId, starTreeDocument.dimensions[idx], dim.getField()); + } + idx++; } } @@ -365,11 +497,17 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A FieldValueConverter aggregatedValueType = metricAggregatorInfos.get(i).getValueAggregators().getAggregatedValueType(); if (aggregatedValueType.equals(LONG)) { if (starTreeDocument.metrics[i] != null) { - metricWriters.get(i).addValue(docId, (long) starTreeDocument.metrics[i]); + ((SortedNumericDocValuesWriterWrapper) (metricWriters.get(i))).addValue( + docId, + (long) starTreeDocument.metrics[i] + ); } } else if (aggregatedValueType.equals(DOUBLE)) { if (starTreeDocument.metrics[i] != null) { - metricWriters.get(i).addValue(docId, NumericUtils.doubleToSortableLong((Double) starTreeDocument.metrics[i])); + ((SortedNumericDocValuesWriterWrapper) (metricWriters.get(i))).addValue( + docId, + NumericUtils.doubleToSortableLong((Double) starTreeDocument.metrics[i]) + ); } } else { throw new IllegalStateException("Unknown metric doc value type"); @@ -379,26 +517,58 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A } } } - addStarTreeDocValueFields(docValuesConsumer, dimensionWriters, dimensionFieldInfoList, numDimensions); addStarTreeDocValueFields(docValuesConsumer, metricWriters, metricFieldInfoList, metricAggregatorInfos.size()); } + /** + * Adds startree field to respective field writers + */ + private void indexDocValue(DocValuesWriterWrapper dvWriter, int docId, long value, String field) throws IOException { + if (dvWriter instanceof SortedSetDocValuesWriterWrapper) { + // TODO : cache lookupOrd to make it faster + ((SortedSetDocValuesWriterWrapper) dvWriter).addValue(docId, sortedSetDocValuesMap.get(field).lookupOrd(value)); + } else if (dvWriter instanceof SortedNumericDocValuesWriterWrapper) { + ((SortedNumericDocValuesWriterWrapper) dvWriter).addValue(docId, value); + } + } + + @SuppressWarnings("unchecked") private void addStarTreeDocValueFields( DocValuesConsumer docValuesConsumer, - List docValuesWriters, + List> docValuesWriters, FieldInfo[] fieldInfoList, int fieldCount ) throws IOException { for (int i = 0; i < fieldCount; i++) { final int writerIndex = i; - DocValuesProducer docValuesProducer = new EmptyDocValuesProducer() { - @Override - public SortedNumericDocValues getSortedNumeric(FieldInfo field) { - return docValuesWriters.get(writerIndex).getDocValues(); - } - }; - docValuesConsumer.addSortedNumericField(fieldInfoList[i], docValuesProducer); + DocValuesProducer docValuesProducer; + switch (fieldInfoList[i].getDocValuesType()) { + case SORTED_NUMERIC: + docValuesProducer = new EmptyDocValuesProducer() { + @Override + public SortedNumericDocValues getSortedNumeric(FieldInfo field) { + DocValuesWriterWrapper wrapper = (DocValuesWriterWrapper< + SortedNumericDocValues>) docValuesWriters.get(writerIndex); + return wrapper.getDocValues(); + } + }; + docValuesConsumer.addSortedNumericField(fieldInfoList[i], docValuesProducer); + break; + case SORTED_SET: + docValuesProducer = new EmptyDocValuesProducer() { + @Override + public SortedSetDocValues getSortedSet(FieldInfo field) { + DocValuesWriterWrapper wrapper = (DocValuesWriterWrapper< + SortedSetDocValues>) docValuesWriters.get(writerIndex); + return wrapper.getDocValues(); + } + }; + docValuesConsumer.addSortedSetField(fieldInfoList[i], docValuesProducer); + break; + default: + throw new IllegalStateException("Unsupported doc values type"); + } } } @@ -408,13 +578,14 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) { protected StarTreeDocument getStarTreeDocument( int currentDocId, SequentialDocValuesIterator[] dimensionReaders, - List metricReaders + List metricReaders, + LongValues longValues ) throws IOException { Long[] dims = new Long[numDimensions]; int i = 0; for (SequentialDocValuesIterator dimensionValueIterator : dimensionReaders) { dimensionValueIterator.nextEntry(currentDocId); - Long val = dimensionValueIterator.value(currentDocId); + Long val = dimensionValueIterator.value(currentDocId, longValues); dims[i] = val; i++; } @@ -455,7 +626,6 @@ protected void setReadersAndNumSegmentDocs( metricReaders.add(new SequentialDocValuesIterator(starTreeValues.getMetricValuesIterator(metricFullName))); } } - numSegmentDocs.set( Integer.parseInt(starTreeValues.getAttributes().getOrDefault(SEGMENT_DOCS_COUNT, String.valueOf(DocIdSetIterator.NO_MORE_DOCS))) ); @@ -664,6 +834,14 @@ private static Long getLong(Object metric) { return metricValue; } + /** + * Sets the sortedSetDocValuesMap. + * This is needed as we need to write the ordinals and also the bytesRef associated with it + */ + protected void setSortedSetDocValuesMap(Map sortedSetDocValuesMap) { + this.sortedSetDocValuesMap = sortedSetDocValuesMap; + } + /** * Merges a star-tree document into an aggregated star-tree document. * A new aggregated star-tree document is created if the aggregated document is null. @@ -730,7 +908,7 @@ private SequentialDocValuesIterator getIteratorForNumericField( * @throws IOException throws an exception if we are unable to add the doc */ private void appendToStarTree(StarTreeDocument starTreeDocument) throws IOException { - + // System.out.println(starTreeDocument); appendStarTreeDocument(starTreeDocument); numStarTreeDocs++; } @@ -795,7 +973,6 @@ private void constructStarTree(InMemoryTreeNode node, int startDocId, int endDoc constructStarTree(child, child.getStartDocId(), child.getEndDocId()); } } - } /** @@ -833,7 +1010,6 @@ private void addChildNode(InMemoryTreeNode node, int endDocId, int dimensionId, childNodeDimensionValue = nodeDimensionValue; childNodeType = StarTreeNodeType.DEFAULT.getValue(); } - InMemoryTreeNode lastNode = getNewNode(dimensionId, nodeStartDocId, endDocId, childNodeType, childNodeDimensionValue); node.addChildNode(lastNode, nodeDimensionValue); } @@ -936,7 +1112,7 @@ public void close() throws IOException { } - abstract Iterator mergeStarTrees(List starTreeValues) throws IOException; + abstract Iterator mergeStarTrees(List starTreeValues, MergeState mergeState) throws IOException; public InMemoryTreeNode getRootNode() { return rootNode; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java index 9c39555396397..f4eee73fca7c4 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java @@ -11,8 +11,12 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.OrdinalMap; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.LongValues; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; @@ -29,6 +33,7 @@ import java.util.Collections; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; @@ -85,11 +90,14 @@ public void appendStarTreeDocument(StarTreeDocument starTreeDocument) throws IOE public void build( List starTreeValuesSubs, AtomicInteger fieldNumberAcrossStarTrees, - DocValuesConsumer starTreeDocValuesConsumer + DocValuesConsumer starTreeDocValuesConsumer, + MergeState mergeState, + Map sortedSetDocValuesMap ) throws IOException { boolean success = false; + setSortedSetDocValuesMap(sortedSetDocValuesMap); try { - build(mergeStarTrees(starTreeValuesSubs), fieldNumberAcrossStarTrees, starTreeDocValuesConsumer); + build(mergeStarTrees(starTreeValuesSubs, mergeState), fieldNumberAcrossStarTrees, starTreeDocValuesConsumer); success = true; } finally { starTreeDocumentFileManager.deleteFiles(success); @@ -135,22 +143,29 @@ public Iterator sortAndAggregateSegmentDocuments( * @param starTreeValuesSubs StarTreeValues from multiple segments * @return iterator of star tree documents */ - Iterator mergeStarTrees(List starTreeValuesSubs) throws IOException { + Iterator mergeStarTrees(List starTreeValuesSubs, MergeState mergeState) throws IOException { int numDocs = 0; int[] docIds; + OrdinalMap ordinalMap = getOrdinalMap(starTreeValuesSubs, mergeState); try { + int seg = 0; for (StarTreeValues starTreeValues : starTreeValuesSubs) { SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[numDimensions]; List metricReaders = new ArrayList<>(); AtomicInteger numSegmentDocs = new AtomicInteger(); setReadersAndNumSegmentDocs(dimensionReaders, metricReaders, numSegmentDocs, starTreeValues); int currentDocId = 0; + LongValues longValues = null; + if (ordinalMap != null) { + longValues = ordinalMap.getGlobalOrds(seg); + } while (currentDocId < numSegmentDocs.get()) { - StarTreeDocument starTreeDocument = getStarTreeDocument(currentDocId, dimensionReaders, metricReaders); + StarTreeDocument starTreeDocument = getStarTreeDocument(currentDocId, dimensionReaders, metricReaders, longValues); segmentDocumentFileManager.writeStarTreeDocument(starTreeDocument, true); numDocs++; currentDocId++; } + seg++; } docIds = new int[numDocs]; for (int i = 0; i < numDocs; i++) { 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 07142fc5c8be7..a6d987974a15b 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 @@ -8,8 +8,12 @@ package org.opensearch.index.compositeindex.datacube.startree.builder; import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.OrdinalMap; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.LongValues; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; @@ -20,8 +24,10 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; @@ -99,9 +105,12 @@ public Iterator sortAndAggregateSegmentDocuments( public void build( List starTreeValuesSubs, AtomicInteger fieldNumberAcrossStarTrees, - DocValuesConsumer starTreeDocValuesConsumer + DocValuesConsumer starTreeDocValuesConsumer, + MergeState mergeState, + Map sortedSetDocValuesMap ) throws IOException { - build(mergeStarTrees(starTreeValuesSubs), fieldNumberAcrossStarTrees, starTreeDocValuesConsumer); + setSortedSetDocValuesMap(sortedSetDocValuesMap); + build(mergeStarTrees(starTreeValuesSubs, mergeState), fieldNumberAcrossStarTrees, starTreeDocValuesConsumer); } /** @@ -112,8 +121,8 @@ public void build( * @return iterator of star tree documents */ @Override - Iterator mergeStarTrees(List starTreeValuesSubs) throws IOException { - return sortAndAggregateStarTreeDocuments(getSegmentsStarTreeDocuments(starTreeValuesSubs), true); + Iterator mergeStarTrees(List starTreeValuesSubs, MergeState mergeState) throws IOException { + return sortAndAggregateStarTreeDocuments(getSegmentsStarTreeDocuments(starTreeValuesSubs, mergeState), true); } /** @@ -123,19 +132,26 @@ Iterator mergeStarTrees(List starTreeValuesSub * @param starTreeValuesSubs StarTreeValues from multiple segments * @return array of star tree documents */ - StarTreeDocument[] getSegmentsStarTreeDocuments(List starTreeValuesSubs) throws IOException { + StarTreeDocument[] getSegmentsStarTreeDocuments(List starTreeValuesSubs, MergeState mergeState) throws IOException { List starTreeDocuments = new ArrayList<>(); + OrdinalMap ordinalMap = getOrdinalMap(starTreeValuesSubs, mergeState); + int seg = 0; for (StarTreeValues starTreeValues : starTreeValuesSubs) { - + Map segToGlobalOrdMap = new HashMap<>(); SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[numDimensions]; List metricReaders = new ArrayList<>(); AtomicInteger numSegmentDocs = new AtomicInteger(); setReadersAndNumSegmentDocs(dimensionReaders, metricReaders, numSegmentDocs, starTreeValues); int currentDocId = 0; + LongValues longValues = null; + if (ordinalMap != null) { + longValues = ordinalMap.getGlobalOrds(seg); + } while (currentDocId < numSegmentDocs.get()) { - starTreeDocuments.add(getStarTreeDocument(currentDocId, dimensionReaders, metricReaders)); + starTreeDocuments.add(getStarTreeDocument(currentDocId, dimensionReaders, metricReaders, longValues)); currentDocId++; } + seg++; } StarTreeDocument[] starTreeDocumentsArr = new StarTreeDocument[starTreeDocuments.size()]; return starTreeDocuments.toArray(starTreeDocumentsArr); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilder.java index 23415ddf29132..6df8d1259835d 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilder.java @@ -10,6 +10,8 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SortedSetDocValues; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; @@ -52,6 +54,8 @@ void build( void build( List starTreeValuesSubs, AtomicInteger fieldNumberAcrossStarTrees, - DocValuesConsumer starTreeDocValuesConsumer + DocValuesConsumer starTreeDocValuesConsumer, + MergeState mergeState, + Map fieldDocIdSetIteratorMap ) throws IOException; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilder.java index bc598c9aeab7c..f975d0c8888c0 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilder.java @@ -12,7 +12,9 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.store.IndexOutput; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; @@ -106,16 +108,19 @@ public void close() throws IOException { /** * Merges star tree fields from multiple segments * - * @param metaOut an IndexInput for star-tree metadata - * @param dataOut an IndexInput for star-tree data - * @param starTreeValuesSubsPerField starTreeValuesSubs per field - * @param starTreeDocValuesConsumer a consumer to write star-tree doc values + * @param metaOut an IndexInput for star-tree metadata + * @param dataOut an IndexInput for star-tree data + * @param starTreeValuesSubsPerField starTreeValuesSubs per field + * @param starTreeDocValuesConsumer a consumer to write star-tree doc values + * @param fieldDocIdSetIteratorMap */ public void buildDuringMerge( IndexOutput metaOut, IndexOutput dataOut, final Map> starTreeValuesSubsPerField, - DocValuesConsumer starTreeDocValuesConsumer + DocValuesConsumer starTreeDocValuesConsumer, + MergeState mergeState, + Map fieldDocIdSetIteratorMap ) throws IOException { logger.debug("Starting merge of {} star-trees with star-tree fields", starTreeValuesSubsPerField.size()); long startTime = System.currentTimeMillis(); @@ -127,10 +132,16 @@ public void buildDuringMerge( } StarTreeField starTreeField = starTreeValuesList.get(0).getStarTreeField(); try (StarTreeBuilder builder = getStarTreeBuilder(metaOut, dataOut, starTreeField, state, mapperService)) { - builder.build(starTreeValuesList, fieldNumberAcrossStarTrees, starTreeDocValuesConsumer); + builder.build( + starTreeValuesList, + fieldNumberAcrossStarTrees, + starTreeDocValuesConsumer, + mergeState, + fieldDocIdSetIteratorMap + ); } } - logger.debug( + logger.info( "Took {} ms to merge {} star-trees with star-tree fields", System.currentTimeMillis() - startTime, starTreeValuesSubsPerField.size() 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 7352c215ee390..57e47b1a5b9d9 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 @@ -10,6 +10,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.store.IndexInput; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.CompositeIndexMetadata; @@ -62,9 +63,10 @@ public class StarTreeMetadata extends CompositeIndexMetadata { private final String starTreeFieldType; /** - * List of dimension fields used in the star-tree. + * Map of dimension fields to their associated DocValuesType.Insertion order needs to be maintained + * as it dictates dimensionSplitOrder */ - private final List dimensionFields; + LinkedHashMap dimensionFieldsToDocValuesMap; /** * List of metrics, containing field names and associated metric statistics. @@ -128,7 +130,7 @@ public StarTreeMetadata( this.starTreeFieldType = this.getCompositeFieldType().getName(); this.version = version; this.numberOfNodes = readNumberOfNodes(); - this.dimensionFields = readStarTreeDimensions(); + this.dimensionFieldsToDocValuesMap = readStarTreeDimensions(); this.metrics = readMetricEntries(); this.segmentAggregatedDocCount = readSegmentAggregatedDocCount(); this.starTreeDocCount = readStarTreeDocCount(); @@ -151,7 +153,7 @@ public StarTreeMetadata( * @param compositeFieldName name of the composite field. Here, name of the star-tree field. * @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 dimensionFieldsToDocValuesMap map of dimensionFields to docValues * @param metrics list of metric entries * @param segmentAggregatedDocCount segment aggregated doc count * @param starTreeDocCount the total number of star tree documents for the segment @@ -167,7 +169,7 @@ public StarTreeMetadata( IndexInput meta, Integer version, Integer numberOfNodes, - List dimensionFields, + LinkedHashMap dimensionFieldsToDocValuesMap, List metrics, Integer segmentAggregatedDocCount, Integer starTreeDocCount, @@ -183,7 +185,7 @@ public StarTreeMetadata( this.starTreeFieldType = compositeFieldType.getName(); this.version = version; this.numberOfNodes = numberOfNodes; - this.dimensionFields = dimensionFields; + this.dimensionFieldsToDocValuesMap = dimensionFieldsToDocValuesMap; this.metrics = metrics; this.segmentAggregatedDocCount = segmentAggregatedDocCount; this.starTreeDocCount = starTreeDocCount; @@ -202,15 +204,14 @@ private int readDimensionsCount() throws IOException { return meta.readVInt(); } - private List readStarTreeDimensions() throws IOException { + private LinkedHashMap readStarTreeDimensions() throws IOException { int dimensionCount = readDimensionsCount(); - List dimensionFields = new ArrayList<>(); + LinkedHashMap dimensionFieldsToDocValuesMap = new LinkedHashMap<>(); for (int i = 0; i < dimensionCount; i++) { - dimensionFields.add(meta.readString()); + dimensionFieldsToDocValuesMap.put(meta.readString(), getDocValuesType(meta, meta.readByte())); } - - return dimensionFields; + return dimensionFieldsToDocValuesMap; } private int readMetricsCount() throws IOException { @@ -314,8 +315,8 @@ public String getStarTreeFieldType() { * * @return star-tree dimension field numbers */ - public List getDimensionFields() { - return dimensionFields; + public Map getDimensionFields() { + return dimensionFieldsToDocValuesMap; } /** @@ -405,4 +406,23 @@ public int getVersion() { public int getNumberOfNodes() { return numberOfNodes; } + + private static DocValuesType getDocValuesType(IndexInput input, byte b) throws IOException { + switch (b) { + case 0: + return DocValuesType.NONE; + case 1: + return DocValuesType.NUMERIC; + case 2: + return DocValuesType.BINARY; + case 3: + return DocValuesType.SORTED; + case 4: + return DocValuesType.SORTED_SET; + case 5: + return DocValuesType.SORTED_NUMERIC; + default: + throw new CorruptIndexException("invalid docvalues byte: " + b, input); + } + } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java index 42e6f3c59866a..569692ce18893 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java @@ -10,6 +10,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.store.IndexOutput; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo; @@ -130,8 +131,9 @@ private static void writeMeta( metaOut.writeVInt(starTreeField.getDimensionNames().size()); // dimensions - for (String dim : starTreeField.getDimensionNames()) { - metaOut.writeString(dim); + for (int i = 0; i < starTreeField.getDimensionNames().size(); i++) { + metaOut.writeString(starTreeField.getDimensionNames().get(i)); + metaOut.writeByte(docValuesByte(starTreeField.getDimensionDocValueTypes().get(i))); } // number of metrics @@ -171,4 +173,24 @@ private static void writeMeta( metaOut.writeVLong(dataFileLength); } + + private static byte docValuesByte(DocValuesType type) { + switch (type) { + case NONE: + return 0; + case NUMERIC: + return 1; + case BINARY: + return 2; + case SORTED: + return 3; + case SORTED_SET: + return 4; + case SORTED_NUMERIC: + return 5; + default: + // BUG + throw new AssertionError("unhandled DocValuesType: " + type); + } + } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java index 255ad343cde32..60e5ae0670495 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java @@ -9,11 +9,14 @@ package org.opensearch.index.compositeindex.datacube.startree.index; import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.store.IndexInput; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesReader; import org.opensearch.index.compositeindex.CompositeIndexMetadata; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.Metric; @@ -25,6 +28,7 @@ import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeFactory; import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import java.io.IOException; @@ -128,8 +132,15 @@ public StarTreeValues( // build dimensions List readDimensions = new ArrayList<>(); - for (String dimension : starTreeMetadata.getDimensionFields()) { - readDimensions.add(new ReadDimension(dimension)); + for (String dimension : starTreeMetadata.getDimensionFields().keySet()) { + readDimensions.add( + new ReadDimension( + dimension, + readState.fieldInfos.fieldInfo( + fullyQualifiedFieldNameForStarTreeDimensionsDocValues(starTreeMetadata.getCompositeFieldName(), dimension) + ).getDocValuesType() + ) + ); } // star-tree field @@ -151,19 +162,27 @@ public StarTreeValues( metricValuesIteratorMap = new LinkedHashMap<>(); // get doc id set iterators for dimensions - for (String dimension : starTreeMetadata.getDimensionFields()) { + for (String dimension : starTreeMetadata.getDimensionFields().keySet()) { dimensionValuesIteratorMap.put(dimension, () -> { try { - SortedNumericDocValues dimensionSortedNumericDocValues = null; + FieldInfo dimensionfieldInfo = null; if (readState != null) { - FieldInfo dimensionfieldInfo = readState.fieldInfos.fieldInfo( + dimensionfieldInfo = readState.fieldInfos.fieldInfo( fullyQualifiedFieldNameForStarTreeDimensionsDocValues(starTreeField.getName(), dimension) ); - if (dimensionfieldInfo != null) { - dimensionSortedNumericDocValues = compositeDocValuesProducer.getSortedNumeric(dimensionfieldInfo); - } } - return new SortedNumericStarTreeValuesIterator(getSortedNumericDocValues(dimensionSortedNumericDocValues)); + assert dimensionfieldInfo != null; + if (dimensionfieldInfo.getDocValuesType().equals(DocValuesType.SORTED_SET)) { + SortedSetDocValues dimensionSortedSetDocValues = compositeDocValuesProducer.getSortedSet(dimensionfieldInfo); + return new SortedSetStarTreeValuesIterator( + Composite99DocValuesReader.getSortedSetDocValues(dimensionSortedSetDocValues) + ); + } else { + SortedNumericDocValues dimensionSortedNumericDocValues = compositeDocValuesProducer.getSortedNumeric( + dimensionfieldInfo + ); + return new SortedNumericStarTreeValuesIterator(getSortedNumericDocValues(dimensionSortedNumericDocValues)); + } } catch (IOException e) { throw new RuntimeException("Error loading dimension StarTreeValuesIterator", e); } 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 9029a451ca4d9..c4d3526648cac 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 @@ -9,8 +9,10 @@ package org.opensearch.index.compositeindex.datacube.startree.utils; +import org.apache.lucene.util.LongValues; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedSetStarTreeValuesIterator; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import java.io.IOException; @@ -81,6 +83,33 @@ public Long value(int currentEntryId) throws IOException { } return ((SortedNumericStarTreeValuesIterator) starTreeValuesIterator).nextValue(); + } else if (starTreeValuesIterator instanceof SortedSetStarTreeValuesIterator) { + if (currentEntryId < 0) { + throw new IllegalStateException("invalid entry id to fetch the next value"); + } + if (currentEntryId == StarTreeValuesIterator.NO_MORE_ENTRIES) { + throw new IllegalStateException("StarTreeValuesIterator is already exhausted"); + } + if (entryId == StarTreeValuesIterator.NO_MORE_ENTRIES || entryId != currentEntryId) { + return null; + } + return ((SortedSetStarTreeValuesIterator) starTreeValuesIterator).nextOrd(); + } else { + throw new IllegalStateException("Unsupported Iterator requested for SequentialDocValuesIterator"); + } + } + + public Long value(int currentEntryId, LongValues globalOrdinalLongValues) throws IOException { + if (starTreeValuesIterator instanceof SortedNumericStarTreeValuesIterator) { + return value(currentEntryId); + } else if (starTreeValuesIterator instanceof SortedSetStarTreeValuesIterator) { + assert globalOrdinalLongValues != null; + Long val = value(currentEntryId); + // convert local ordinal to global ordinal + if (val != null) { + val = globalOrdinalLongValues.get(val); + } + return val; } else { throw new IllegalStateException("Unsupported Iterator requested for SequentialDocValuesIterator"); } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java index 2aae0d4ca7e29..240a727678d6f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java @@ -16,6 +16,7 @@ import java.util.Collections; import java.util.List; +import java.util.Map; /** * Util class for building star tree @@ -67,14 +68,17 @@ public static String fullyQualifiedFieldNameForStarTreeMetricsDocValues(String s * @param fields field names * @return field infos */ - public static FieldInfo[] getFieldInfoList(List fields) { + public static FieldInfo[] getFieldInfoList(List fields, Map dimDocValuesTypeMap) { FieldInfo[] fieldInfoList = new FieldInfo[fields.size()]; - // field number is not really used. We depend on unique field names to get the desired iterator int fieldNumber = 0; - for (String fieldName : fields) { - fieldInfoList[fieldNumber] = getFieldInfo(fieldName, DocValuesType.SORTED_NUMERIC, fieldNumber); + fieldInfoList[fieldNumber] = getFieldInfo( + fieldName, + // default is sortedNumeric since all metrics right now are sorted numeric + dimDocValuesTypeMap.getOrDefault(fieldName, DocValuesType.SORTED_NUMERIC), + fieldNumber + ); fieldNumber++; } return fieldInfoList; diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java new file mode 100644 index 0000000000000..0cddffe5877e9 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java @@ -0,0 +1,59 @@ +/* + * 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.iterator; + +import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.automaton.CompiledAutomaton; +import org.opensearch.common.annotation.ExperimentalApi; + +import java.io.IOException; + +/** + * Wrapper iterator class for StarTree index to traverse through SortedNumericDocValues + * + * @opensearch.experimental + */ +@ExperimentalApi +public class SortedSetStarTreeValuesIterator extends StarTreeValuesIterator { + + public SortedSetStarTreeValuesIterator(DocIdSetIterator docIdSetIterator) { + super(docIdSetIterator); + } + + public long nextOrd() throws IOException { + return ((SortedSetDocValues) docIdSetIterator).nextOrd(); + } + + public int docValueCount() { + return ((SortedSetDocValues) docIdSetIterator).docValueCount(); + } + + public BytesRef lookupOrd(long ord) throws IOException { + return ((SortedSetDocValues) docIdSetIterator).lookupOrd(ord); + } + + public long getValueCount() { + return ((SortedSetDocValues) docIdSetIterator).getValueCount(); + } + + public long lookupTerm(BytesRef key) throws IOException { + return ((SortedSetDocValues) docIdSetIterator).lookupTerm(key); + } + + public TermsEnum termsEnum() throws IOException { + return ((SortedSetDocValues) docIdSetIterator).termsEnum(); + } + + public TermsEnum intersect(CompiledAutomaton automaton) throws IOException { + return ((SortedSetDocValues) docIdSetIterator).intersect(automaton); + } +} diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java index 4fe0199f89f41..158c9dad583bc 100644 --- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java +++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java @@ -118,24 +118,24 @@ public void testStarTreeDocValues() throws IOException { RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf); Document doc = new Document(); doc.add(new SortedNumericDocValuesField("sndv", 1)); - doc.add(new SortedNumericDocValuesField("dv", 1)); - doc.add(new SortedNumericDocValuesField("field", -1)); + doc.add(new SortedNumericDocValuesField("dv1", 1)); + doc.add(new SortedNumericDocValuesField("field1", -1)); iw.addDocument(doc); doc = new Document(); doc.add(new SortedNumericDocValuesField("sndv", 1)); - doc.add(new SortedNumericDocValuesField("dv", 1)); - doc.add(new SortedNumericDocValuesField("field", -1)); + doc.add(new SortedNumericDocValuesField("dv1", 1)); + doc.add(new SortedNumericDocValuesField("field1", -1)); iw.addDocument(doc); doc = new Document(); iw.forceMerge(1); doc.add(new SortedNumericDocValuesField("sndv", 2)); - doc.add(new SortedNumericDocValuesField("dv", 2)); - doc.add(new SortedNumericDocValuesField("field", -2)); + doc.add(new SortedNumericDocValuesField("dv1", 2)); + doc.add(new SortedNumericDocValuesField("field1", -2)); iw.addDocument(doc); doc = new Document(); doc.add(new SortedNumericDocValuesField("sndv", 2)); - doc.add(new SortedNumericDocValuesField("dv", 2)); - doc.add(new SortedNumericDocValuesField("field", -2)); + doc.add(new SortedNumericDocValuesField("dv1", 2)); + doc.add(new SortedNumericDocValuesField("field1", -2)); iw.addDocument(doc); iw.forceMerge(1); iw.close(); @@ -219,12 +219,12 @@ private XContentBuilder getExpandedMapping() throws IOException { b.field("name", "sndv"); b.endObject(); b.startObject(); - b.field("name", "dv"); + b.field("name", "dv1"); b.endObject(); b.endArray(); b.startArray("metrics"); b.startObject(); - b.field("name", "field"); + b.field("name", "field1"); b.startArray("stats"); b.value("sum"); b.value("value_count"); @@ -251,10 +251,10 @@ private XContentBuilder getExpandedMapping() throws IOException { b.startObject("sndv"); b.field("type", "integer"); b.endObject(); - b.startObject("dv"); + b.startObject("dv1"); b.field("type", "integer"); b.endObject(); - b.startObject("field"); + b.startObject("field1"); b.field("type", "integer"); b.endObject(); b.endObject(); diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java index 705700aada2ff..bd80d8f39d86f 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java @@ -16,8 +16,10 @@ import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; @@ -183,7 +185,9 @@ public static void setup() throws IOException { public void build( List starTreeValuesSubs, AtomicInteger fieldNumberAcrossStarTrees, - DocValuesConsumer starTreeDocValuesConsumer + DocValuesConsumer starTreeDocValuesConsumer, + MergeState mergeState, + Map fieldDocIdSetIteratorMap ) throws IOException {} @Override @@ -219,7 +223,7 @@ public Iterator generateStarTreeDocumentsForStarNode(int start } @Override - Iterator mergeStarTrees(List starTreeValues) throws IOException { + Iterator mergeStarTrees(List starTreeValues, MergeState mergeState) throws IOException { return null; } }; diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java index bb31bd6a7cc27..e7fae41c0f67f 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java @@ -386,7 +386,7 @@ public static void validateStarTreeFileFormats( public static SegmentReadState getReadState( int numDocs, - List dimensionFields, + Map dimensionFields, List metrics, StarTreeField compositeField, SegmentWriteState writeState, @@ -401,7 +401,7 @@ public static SegmentReadState getReadState( FieldInfo[] fields = new FieldInfo[dimensionFields.size() + numMetrics]; int i = 0; - for (String dimension : dimensionFields) { + for (String dimension : dimensionFields.keySet()) { fields[i] = new FieldInfo( fullyQualifiedFieldNameForStarTreeDimensionsDocValues(compositeField.getName(), dimension), i, diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java index 095eda2986b3a..a46909d696bf6 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java @@ -52,6 +52,7 @@ import java.util.Collections; import java.util.HashMap; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -591,13 +592,19 @@ public void test_build_multipleStarTrees() throws IOException { metaOut.close(); dataOut.close(); + LinkedHashMap fieldsMap = new LinkedHashMap<>(); + fieldsMap.put("field1", DocValuesType.SORTED_NUMERIC); + fieldsMap.put("field3", DocValuesType.SORTED_NUMERIC); + fieldsMap.put("field5", DocValuesType.SORTED_NUMERIC); + fieldsMap.put("field8", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( "test", STAR_TREE, mock(IndexInput.class), VERSION_CURRENT, builder.numStarTreeNodes, - List.of("field1", "field3", "field5", "field8"), + fieldsMap, List.of( new Metric("field2", List.of(MetricStat.SUM)), new Metric("field4", List.of(MetricStat.SUM)), @@ -614,13 +621,18 @@ public void test_build_multipleStarTrees() throws IOException { 330 ); + LinkedHashMap fieldsMap1 = new LinkedHashMap<>(); + fieldsMap1.put("fieldC", DocValuesType.SORTED_NUMERIC); + fieldsMap1.put("fieldB", DocValuesType.SORTED_NUMERIC); + fieldsMap1.put("fieldL", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata2 = new StarTreeMetadata( "test", STAR_TREE, mock(IndexInput.class), VERSION_CURRENT, builder.numStarTreeNodes, - List.of("fieldC", "fieldB", "fieldL"), + fieldsMap1, List.of(new Metric("fieldI", List.of(MetricStat.SUM))), 7, 27, @@ -631,9 +643,8 @@ public void test_build_multipleStarTrees() throws IOException { 1287 ); - List totalDimensionFields = new ArrayList<>(); - totalDimensionFields.addAll(starTreeMetadata.getDimensionFields()); - totalDimensionFields.addAll(starTreeMetadata2.getDimensionFields()); + LinkedHashMap totalDimensionFields = new LinkedHashMap<>(starTreeMetadata.getDimensionFields()); + totalDimensionFields.putAll(starTreeMetadata2.getDimensionFields()); List metrics = new ArrayList<>(); metrics.addAll(starTreeMetadata.getMetrics()); diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java index 1aa830e3587df..f232d8fb9a686 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java @@ -9,6 +9,7 @@ package org.opensearch.index.compositeindex.datacube.startree.builder; import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.EmptyDocValuesProducer; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.SegmentWriteState; @@ -32,6 +33,7 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -124,14 +126,16 @@ public void testFlushFlow() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - + LinkedHashMap docValues = new LinkedHashMap<>(); + docValues.put("field1", DocValuesType.SORTED_NUMERIC); + docValues.put("field3", DocValuesType.SORTED_NUMERIC); StarTreeMetadata starTreeMetadata = new StarTreeMetadata( "sf", STAR_TREE, mock(IndexInput.class), VERSION_CURRENT, builder.numStarTreeNodes, - List.of("field1", "field3"), + docValues, List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), 6, builder.numStarTreeDocs, @@ -222,13 +226,16 @@ public void testFlushFlowDimsReverse() throws IOException { dataOut.close(); docValuesConsumer.close(); + LinkedHashMap docValues = new LinkedHashMap<>(); + docValues.put("field1", DocValuesType.SORTED_NUMERIC); + docValues.put("field3", DocValuesType.SORTED_NUMERIC); StarTreeMetadata starTreeMetadata = new StarTreeMetadata( "sf", STAR_TREE, mock(IndexInput.class), VERSION_CURRENT, builder.numStarTreeNodes, - List.of("field1", "field3"), + docValues, List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), 6, builder.numStarTreeDocs, @@ -322,7 +329,10 @@ public void testFlushFlowBuild() throws IOException { dataOut.close(); docValuesConsumer.close(); - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 100, 1, 6699); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 100, 1, 6699); validateStarTreeFileFormats( builder.getRootNode(), diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java index f983365dfec30..6d034cc91700a 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java @@ -9,6 +9,7 @@ package org.opensearch.index.compositeindex.datacube.startree.builder; import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.DocIdSetIterator; @@ -200,7 +201,7 @@ public void testMergeFlow() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** [0, 0, 0, 0] | [0.0, 2] [1, 1, 1, 1] | [20.0, 2] @@ -316,7 +317,7 @@ public void testMergeFlowWithSum() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Sum [ metric] ] * [0, 0] | [0.0] @@ -347,8 +348,10 @@ public void testMergeFlowWithSum() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 6, 1000, 264); validateStarTreeFileFormats( builder.getRootNode(), @@ -392,7 +395,7 @@ public void testMergeFlowWithCount() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -420,8 +423,10 @@ public void testMergeFlowWithCount() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 6, 1000, 264); validateStarTreeFileFormats( builder.getRootNode(), @@ -466,7 +471,7 @@ public void testMergeFlowNumSegmentsDocs() throws IOException { "4" ); builder = getStarTreeBuilder(metaOut, dataOut, sf, getWriteState(4, writeState.segmentInfo.getId()), mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -534,7 +539,7 @@ public void testMergeFlowWithMissingDocs() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -567,8 +572,10 @@ public void testMergeFlowWithMissingDocs() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 10, 1000, 363); validateStarTreeFileFormats( builder.getRootNode(), @@ -621,7 +628,7 @@ public void testMergeFlowWithMissingDocsWithZero() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [9] @@ -655,8 +662,10 @@ public void testMergeFlowWithMissingDocsWithZero() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 231); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 6, 1000, 231); validateStarTreeFileFormats( builder.getRootNode(), @@ -708,7 +717,7 @@ public void testMergeFlowWithMissingDocsWithZeroComplexCase() throws IOException Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [9] @@ -746,8 +755,10 @@ public void testMergeFlowWithMissingDocsWithZeroComplexCase() throws IOException metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 7, 1000, 231); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 7, 1000, 231); validateStarTreeFileFormats( builder.getRootNode(), @@ -799,7 +810,7 @@ public void testMergeFlowWithMissingDocsInSecondDim() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -833,8 +844,10 @@ public void testMergeFlowWithMissingDocsInSecondDim() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 10, 1000, 363); validateStarTreeFileFormats( builder.getRootNode(), @@ -885,7 +898,7 @@ public void testMergeFlowWithDocsMissingAtTheEnd() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -918,8 +931,10 @@ public void testMergeFlowWithDocsMissingAtTheEnd() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 10, 1000, 363); validateStarTreeFileFormats( builder.getRootNode(), @@ -963,7 +978,7 @@ public void testMergeFlowWithEmptyFieldsInOneSegment() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -991,8 +1006,10 @@ public void testMergeFlowWithEmptyFieldsInOneSegment() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 6, 1000, 264); validateStarTreeFileFormats( builder.getRootNode(), @@ -1092,7 +1109,7 @@ public void testMergeFlowWithDuplicateDimensionValues() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState), new AtomicInteger(), docValuesConsumer); List starTreeDocuments = builder.getStarTreeDocuments(); assertEquals(401, starTreeDocuments.size()); int count = 0; @@ -1240,7 +1257,7 @@ public void testMergeFlowWithMaxLeafDocs() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState), new AtomicInteger(), docValuesConsumer); List starTreeDocuments = builder.getStarTreeDocuments(); /** 635 docs get generated @@ -1358,7 +1375,7 @@ public void testMergeFlowWithDifferentDocsFromSegments() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] [0, 0] | [0] @@ -1390,8 +1407,10 @@ public void testMergeFlowWithDifferentDocsFromSegments() throws IOException { metaOut.close(); dataOut.close(); docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 9, 1000, 330); + LinkedHashMap map = new LinkedHashMap<>(); + map.put("field1", DocValuesType.SORTED_NUMERIC); + map.put("field3", DocValuesType.SORTED_NUMERIC); + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(map, 9, 1000, 330); validateStarTreeFileFormats( builder.getRootNode(), @@ -1495,7 +1514,7 @@ public void testMergeFlowWithDuplicateDimensionValueWithMaxLeafDocs() throws IOE Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState), new AtomicInteger(), docValuesConsumer); List starTreeDocuments = builder.getStarTreeDocuments(); assertEquals(401, starTreeDocuments.size()); validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); @@ -1624,7 +1643,7 @@ public void testMergeFlowWithMaxLeafDocsAndStarTreeNodesAssertion() throws IOExc Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState), new AtomicInteger(), docValuesConsumer); List starTreeDocuments = builder.getStarTreeDocuments(); Map> dimValueToDocIdMap = new HashMap<>(); traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); @@ -1733,7 +1752,7 @@ public void testMergeFlowWithTimestamps() throws IOException { Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION ); builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(4, writeState.segmentInfo.getId()), mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2), mergeState); /** [1655287972000, 1655287972000, 1655287972000, 3] | [30.0, 3] [1655288032000, 1655288032000, 1655288032000, 2] | [20.0, 2] diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java index 4c854f7546197..0525652dbeb5f 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java @@ -16,6 +16,7 @@ import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.VectorEncoding; @@ -60,6 +61,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -87,6 +89,7 @@ public abstract class StarTreeBuilderTestCase extends OpenSearchTestCase { protected String dataFileName; protected String metaFileName; protected List dimensionsOrder; + protected MergeState mergeState; public StarTreeBuilderTestCase(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { this.buildMode = buildMode; @@ -155,6 +158,8 @@ public void setup() throws IOException { } writeState = getWriteState(5, UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8)); + mergeState = new MergeState(null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, false); + dataFileName = IndexFileNames.segmentFileName( writeState.segmentInfo.name, writeState.segmentSuffix, @@ -240,7 +245,7 @@ SegmentWriteState getWriteState(int numDocs, byte[] id) { return BuilderTestsUtils.getWriteState(numDocs, id, fieldsInfo, directory); } - SegmentReadState getReadState(int numDocs, List dimensionFields, List metrics) { + SegmentReadState getReadState(int numDocs, Map dimensionFields, List metrics) { return BuilderTestsUtils.getReadState(numDocs, dimensionFields, metrics, compositeField, writeState, directory); } @@ -248,10 +253,12 @@ protected Map getAttributes(int numSegmentDocs) { return Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, String.valueOf(numSegmentDocs)); } - protected List getStarTreeDimensionNames(List dimensionsOrder) { - List dimensionNames = new ArrayList<>(); + protected LinkedHashMap getStarTreeDimensionNames(List dimensionsOrder) { + LinkedHashMap dimensionNames = new LinkedHashMap<>(); for (Dimension dimension : dimensionsOrder) { - dimensionNames.addAll(dimension.getSubDimensionNames()); + for (String dimensionName : dimension.getSubDimensionNames()) { + dimensionNames.put(dimensionName, dimension.getDocValuesType()); + } } return dimensionNames; } @@ -320,7 +327,12 @@ protected long getLongFromDouble(double value) { return NumericUtils.doubleToSortableLong(value); } - protected StarTreeMetadata getStarTreeMetadata(List fields, int segmentAggregatedDocCount, int maxLeafDocs, int dataLength) { + protected StarTreeMetadata getStarTreeMetadata( + LinkedHashMap fields, + int segmentAggregatedDocCount, + int maxLeafDocs, + int dataLength + ) { return new StarTreeMetadata( "sf", STAR_TREE, diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java index c8636426449ad..cc91d69be97c1 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java @@ -42,6 +42,7 @@ import java.util.Collections; import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.UUID; @@ -181,8 +182,11 @@ public void test_starTreeMetadata() throws IOException { assertEquals(starTreeMetadata.getNumberOfNodes(), numberOfNodes); assertNotNull(starTreeMetadata); - for (int i = 0; i < dimensionsOrder.size(); i++) { - assertEquals(dimensionsOrder.get(i).getField(), starTreeMetadata.getDimensionFields().get(i)); + assertEquals(dimensionsOrder.size(), starTreeMetadata.dimensionFieldsToDocValuesMap.size()); + int k = 0; + for (Map.Entry entry : starTreeMetadata.dimensionFieldsToDocValuesMap.entrySet()) { + assertEquals(dimensionsOrder.get(k).getField(), entry.getKey()); + k++; } assertEquals(starTreeField.getMetrics().size(), starTreeMetadata.getMetrics().size()); diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtilsTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtilsTests.java index 9cca0b04e9ea4..7e438c18d9ab9 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtilsTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtilsTests.java @@ -17,6 +17,7 @@ import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.UUID; @@ -43,7 +44,7 @@ public void testFullyQualifiedFieldNameForStarTreeMetricsDocValues() { public void testGetFieldInfoList() { List fieldNames = Arrays.asList("field1", "field2", "field3"); - FieldInfo[] actualFieldInfos = StarTreeUtils.getFieldInfoList(fieldNames); + FieldInfo[] actualFieldInfos = StarTreeUtils.getFieldInfoList(fieldNames, new HashMap<>()); for (int i = 0; i < fieldNames.size(); i++) { assertFieldInfos(actualFieldInfos[i], fieldNames.get(i), i); }