Skip to content

Commit 3314b2d

Browse files
opensearch-trigger-bot[bot]github-actions[bot]
opensearch-trigger-bot[bot]
and
github-actions[bot]
authored
[Star tree] Star tree merge changes (#14652) (#15380)
--------- (cherry picked from commit 312de99) Signed-off-by: Bharathwaj G <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 4957818 commit 3314b2d

32 files changed

+3281
-1435
lines changed

‎server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesReader.java‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@
1717
import org.apache.lucene.index.SortedNumericDocValues;
1818
import org.apache.lucene.index.SortedSetDocValues;
1919
import org.opensearch.common.annotation.ExperimentalApi;
20-
import org.opensearch.index.mapper.CompositeMappedFieldType;
2120

2221
import java.io.IOException;
22+
import java.util.ArrayList;
2323
import java.util.List;
2424

2525
/**
@@ -74,15 +74,13 @@ public void close() throws IOException {
7474
}
7575

7676
@Override
77-
public List<String> getCompositeIndexFields() {
77+
public List<CompositeIndexFieldInfo> getCompositeIndexFields() {
7878
// todo : read from file formats and get the field names.
79-
throw new UnsupportedOperationException();
80-
79+
return new ArrayList<>();
8180
}
8281

8382
@Override
84-
public CompositeIndexValues getCompositeIndexValues(String field, CompositeMappedFieldType.CompositeFieldType fieldType)
85-
throws IOException {
83+
public CompositeIndexValues getCompositeIndexValues(CompositeIndexFieldInfo compositeIndexFieldInfo) throws IOException {
8684
// TODO : read compositeIndexValues [starTreeValues] from star tree files
8785
throw new UnsupportedOperationException();
8886
}

‎server/src/main/java/org/opensearch/index/codec/composite/Composite99DocValuesWriter.java‎

Lines changed: 91 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,29 @@
88

99
package org.opensearch.index.codec.composite;
1010

11+
import org.apache.logging.log4j.LogManager;
12+
import org.apache.logging.log4j.Logger;
1113
import org.apache.lucene.codecs.DocValuesConsumer;
1214
import org.apache.lucene.codecs.DocValuesProducer;
15+
import org.apache.lucene.index.DocValues;
16+
import org.apache.lucene.index.DocValuesType;
17+
import org.apache.lucene.index.EmptyDocValuesProducer;
1318
import org.apache.lucene.index.FieldInfo;
1419
import org.apache.lucene.index.MergeState;
1520
import org.apache.lucene.index.SegmentWriteState;
21+
import org.apache.lucene.index.SortedNumericDocValues;
1622
import org.opensearch.common.annotation.ExperimentalApi;
23+
import org.opensearch.index.codec.composite.datacube.startree.StarTreeValues;
24+
import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
1725
import org.opensearch.index.compositeindex.datacube.startree.builder.StarTreesBuilder;
1826
import org.opensearch.index.mapper.CompositeMappedFieldType;
1927
import org.opensearch.index.mapper.MapperService;
20-
import org.opensearch.index.mapper.StarTreeMapper;
2128

2229
import java.io.IOException;
30+
import java.util.Collections;
2331
import java.util.HashMap;
2432
import java.util.HashSet;
33+
import java.util.List;
2534
import java.util.Map;
2635
import java.util.Set;
2736
import java.util.concurrent.atomic.AtomicReference;
@@ -40,8 +49,10 @@ public class Composite99DocValuesWriter extends DocValuesConsumer {
4049
AtomicReference<MergeState> mergeState = new AtomicReference<>();
4150
private final Set<CompositeMappedFieldType> compositeMappedFieldTypes;
4251
private final Set<String> compositeFieldSet;
52+
private final Set<String> segmentFieldSet;
4353

4454
private final Map<String, DocValuesProducer> fieldProducerMap = new HashMap<>();
55+
private static final Logger logger = LogManager.getLogger(Composite99DocValuesWriter.class);
4556

4657
public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState segmentWriteState, MapperService mapperService) {
4758

@@ -50,6 +61,12 @@ public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState
5061
this.mapperService = mapperService;
5162
this.compositeMappedFieldTypes = mapperService.getCompositeFieldTypes();
5263
compositeFieldSet = new HashSet<>();
64+
segmentFieldSet = new HashSet<>();
65+
for (FieldInfo fi : segmentWriteState.fieldInfos) {
66+
if (DocValuesType.SORTED_NUMERIC.equals(fi.getDocValuesType())) {
67+
segmentFieldSet.add(fi.name);
68+
}
69+
}
5370
for (CompositeMappedFieldType type : compositeMappedFieldTypes) {
5471
compositeFieldSet.addAll(type.fields());
5572
}
@@ -95,23 +112,91 @@ private void createCompositeIndicesIfPossible(DocValuesProducer valuesProducer,
95112
fieldProducerMap.put(field.name, valuesProducer);
96113
compositeFieldSet.remove(field.name);
97114
}
115+
segmentFieldSet.remove(field.name);
116+
if (segmentFieldSet.isEmpty()) {
117+
Set<String> compositeFieldSetCopy = new HashSet<>(compositeFieldSet);
118+
for (String compositeField : compositeFieldSetCopy) {
119+
fieldProducerMap.put(compositeField, new EmptyDocValuesProducer() {
120+
@Override
121+
public SortedNumericDocValues getSortedNumeric(FieldInfo field) {
122+
return DocValues.emptySortedNumeric();
123+
}
124+
});
125+
compositeFieldSet.remove(compositeField);
126+
}
127+
}
98128
// we have all the required fields to build composite fields
99129
if (compositeFieldSet.isEmpty()) {
100130
for (CompositeMappedFieldType mappedType : compositeMappedFieldTypes) {
101-
if (mappedType instanceof StarTreeMapper.StarTreeFieldType) {
102-
try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(fieldProducerMap, state, mapperService)) {
103-
starTreesBuilder.build();
131+
if (mappedType.getCompositeIndexType().equals(CompositeMappedFieldType.CompositeFieldType.STAR_TREE)) {
132+
try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService)) {
133+
starTreesBuilder.build(fieldProducerMap);
104134
}
105135
}
106136
}
107137
}
138+
108139
}
109140

110141
@Override
111142
public void merge(MergeState mergeState) throws IOException {
112143
this.mergeState.compareAndSet(null, mergeState);
113144
super.merge(mergeState);
114-
// TODO : handle merge star tree
115-
// mergeStarTreeFields(mergeState);
145+
mergeCompositeFields(mergeState);
146+
}
147+
148+
/**
149+
* Merges composite fields from multiple segments
150+
* @param mergeState merge state
151+
*/
152+
private void mergeCompositeFields(MergeState mergeState) throws IOException {
153+
mergeStarTreeFields(mergeState);
154+
}
155+
156+
/**
157+
* Merges star tree data fields from multiple segments
158+
* @param mergeState merge state
159+
*/
160+
private void mergeStarTreeFields(MergeState mergeState) throws IOException {
161+
Map<String, List<StarTreeValues>> starTreeSubsPerField = new HashMap<>();
162+
StarTreeField starTreeField = null;
163+
for (int i = 0; i < mergeState.docValuesProducers.length; i++) {
164+
CompositeIndexReader reader = null;
165+
if (mergeState.docValuesProducers[i] == null) {
166+
continue;
167+
}
168+
if (mergeState.docValuesProducers[i] instanceof CompositeIndexReader) {
169+
reader = (CompositeIndexReader) mergeState.docValuesProducers[i];
170+
} else {
171+
continue;
172+
}
173+
174+
List<CompositeIndexFieldInfo> compositeFieldInfo = reader.getCompositeIndexFields();
175+
for (CompositeIndexFieldInfo fieldInfo : compositeFieldInfo) {
176+
if (fieldInfo.getType().equals(CompositeMappedFieldType.CompositeFieldType.STAR_TREE)) {
177+
CompositeIndexValues compositeIndexValues = reader.getCompositeIndexValues(fieldInfo);
178+
if (compositeIndexValues instanceof StarTreeValues) {
179+
StarTreeValues starTreeValues = (StarTreeValues) compositeIndexValues;
180+
List<StarTreeValues> fieldsList = starTreeSubsPerField.getOrDefault(fieldInfo.getField(), Collections.emptyList());
181+
if (starTreeField == null) {
182+
starTreeField = starTreeValues.getStarTreeField();
183+
}
184+
// assert star tree configuration is same across segments
185+
else {
186+
if (starTreeField.equals(starTreeValues.getStarTreeField()) == false) {
187+
throw new IllegalArgumentException(
188+
"star tree field configuration must match the configuration of the field being merged"
189+
);
190+
}
191+
}
192+
fieldsList.add(starTreeValues);
193+
starTreeSubsPerField.put(fieldInfo.getField(), fieldsList);
194+
}
195+
}
196+
}
197+
}
198+
try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService)) {
199+
starTreesBuilder.buildDuringMerge(starTreeSubsPerField);
200+
}
116201
}
117202
}

‎server/src/main/java/org/opensearch/index/codec/composite/CompositeIndexFieldInfo.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
2+
/*
3+
* SPDX-License-Identifier: Apache-2.0
4+
*
5+
* The OpenSearch Contributors require contributions made to
6+
* this file be licensed under the Apache-2.0 license or a
7+
* compatible open source license.
8+
*/
9+
10+
package org.opensearch.index.codec.composite;
11+
12+
import org.opensearch.common.annotation.ExperimentalApi;
13+
import org.opensearch.index.mapper.CompositeMappedFieldType;
14+
15+
/**
16+
* Field info details of composite index fields
17+
*
18+
* @opensearch.experimental
19+
*/
20+
@ExperimentalApi
21+
public class CompositeIndexFieldInfo {
22+
private final String field;
23+
private final CompositeMappedFieldType.CompositeFieldType type;
24+
25+
public CompositeIndexFieldInfo(String field, CompositeMappedFieldType.CompositeFieldType type) {
26+
this.field = field;
27+
this.type = type;
28+
}
29+
30+
public String getField() {
31+
return field;
32+
}
33+
34+
public CompositeMappedFieldType.CompositeFieldType getType() {
35+
return type;
36+
}
37+
}

‎server/src/main/java/org/opensearch/index/codec/composite/CompositeIndexReader.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
package org.opensearch.index.codec.composite;
1010

1111
import org.opensearch.common.annotation.ExperimentalApi;
12-
import org.opensearch.index.mapper.CompositeMappedFieldType;
1312

1413
import java.io.IOException;
1514
import java.util.List;
@@ -25,10 +24,10 @@ public interface CompositeIndexReader {
2524
* Get list of composite index fields from the segment
2625
*
2726
*/
28-
List<String> getCompositeIndexFields();
27+
List<CompositeIndexFieldInfo> getCompositeIndexFields();
2928

3029
/**
3130
* Get composite index values based on the field name and the field type
3231
*/
33-
CompositeIndexValues getCompositeIndexValues(String field, CompositeMappedFieldType.CompositeFieldType fieldType) throws IOException;
32+
CompositeIndexValues getCompositeIndexValues(CompositeIndexFieldInfo fieldInfo) throws IOException;
3433
}

‎server/src/main/java/org/opensearch/index/codec/composite/datacube/startree/StarTreeValues.java‎

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88

99
package org.opensearch.index.codec.composite.datacube.startree;
1010

11+
import org.apache.lucene.search.DocIdSetIterator;
1112
import org.opensearch.common.annotation.ExperimentalApi;
1213
import org.opensearch.index.codec.composite.CompositeIndexValues;
14+
import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
15+
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
1316

14-
import java.util.List;
17+
import java.util.Map;
1518

1619
/**
1720
* Concrete class that holds the star tree associated values from the segment
@@ -20,16 +23,48 @@
2023
*/
2124
@ExperimentalApi
2225
public class StarTreeValues implements CompositeIndexValues {
23-
private final List<String> dimensionsOrder;
26+
private final StarTreeField starTreeField;
27+
private final StarTreeNode root;
28+
private final Map<String, DocIdSetIterator> dimensionDocValuesIteratorMap;
29+
private final Map<String, DocIdSetIterator> metricDocValuesIteratorMap;
30+
private final Map<String, String> attributes;
2431

25-
// TODO : come up with full set of vales such as dimensions and metrics doc values + star tree
26-
public StarTreeValues(List<String> dimensionsOrder) {
27-
super();
28-
this.dimensionsOrder = List.copyOf(dimensionsOrder);
32+
public StarTreeValues(
33+
StarTreeField starTreeField,
34+
StarTreeNode root,
35+
Map<String, DocIdSetIterator> dimensionDocValuesIteratorMap,
36+
Map<String, DocIdSetIterator> metricDocValuesIteratorMap,
37+
Map<String, String> attributes
38+
) {
39+
this.starTreeField = starTreeField;
40+
this.root = root;
41+
this.dimensionDocValuesIteratorMap = dimensionDocValuesIteratorMap;
42+
this.metricDocValuesIteratorMap = metricDocValuesIteratorMap;
43+
this.attributes = attributes;
2944
}
3045

3146
@Override
3247
public CompositeIndexValues getValues() {
3348
return this;
3449
}
50+
51+
public StarTreeField getStarTreeField() {
52+
return starTreeField;
53+
}
54+
55+
public StarTreeNode getRoot() {
56+
return root;
57+
}
58+
59+
public Map<String, DocIdSetIterator> getDimensionDocValuesIteratorMap() {
60+
return dimensionDocValuesIteratorMap;
61+
}
62+
63+
public Map<String, DocIdSetIterator> getMetricDocValuesIteratorMap() {
64+
return metricDocValuesIteratorMap;
65+
}
66+
67+
public Map<String, String> getAttributes() {
68+
return attributes;
69+
}
3570
}

‎server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@
1818
public class CountValueAggregator implements ValueAggregator<Long> {
1919
public static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG;
2020
public static final long DEFAULT_INITIAL_VALUE = 1L;
21+
private StarTreeNumericType starTreeNumericType;
22+
23+
public CountValueAggregator(StarTreeNumericType starTreeNumericType) {
24+
this.starTreeNumericType = starTreeNumericType;
25+
}
2126

2227
@Override
2328
public MetricStat getAggregationType() {
@@ -30,12 +35,12 @@ public StarTreeNumericType getAggregatedValueType() {
3035
}
3136

3237
@Override
33-
public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue, StarTreeNumericType starTreeNumericType) {
38+
public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) {
3439
return DEFAULT_INITIAL_VALUE;
3540
}
3641

3742
@Override
38-
public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue, StarTreeNumericType starTreeNumericType) {
43+
public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue) {
3944
return value + 1;
4045
}
4146

@@ -60,7 +65,7 @@ public Long toLongValue(Long value) {
6065
}
6166

6267
@Override
63-
public Long toStarTreeNumericTypeValue(Long value, StarTreeNumericType type) {
68+
public Long toStarTreeNumericTypeValue(Long value) {
6469
return value;
6570
}
6671
}

0 commit comments

Comments
 (0)