Skip to content

Commit 775330e

Browse files
expaniexpani
committed
Increased UT coverage and addressed PR Comments
Signed-off-by: expani <[email protected]>
1 parent 7153ccd commit 775330e

File tree

11 files changed

+266
-104
lines changed

11 files changed

+266
-104
lines changed

‎server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -192,20 +192,6 @@ public StarTreeNode getChildStarNode() throws IOException {
192192
return handleStarNode();
193193
}
194194

195-
@Override
196-
public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
197-
// there will be no children for leaf nodes
198-
if (isLeaf()) {
199-
return null;
200-
}
201-
202-
StarTreeNode resultStarTreeNode = null;
203-
if (null != dimensionValue) {
204-
resultStarTreeNode = binarySearchChild(dimensionValue, null);
205-
}
206-
return resultStarTreeNode;
207-
}
208-
209195
@Override
210196
public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
211197
// there will be no children for leaf nodes

‎server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,16 +108,25 @@ public interface StarTreeNode {
108108
* @return the child node for the given dimension value or null if child is not present
109109
* @throws IOException if an I/O error occurs while retrieving the child node
110110
*/
111-
StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException;
111+
default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
112+
return getChildForDimensionValue(dimensionValue, null);
113+
}
112114

115+
/**
116+
* Matches the given @dimensionValue amongst the child default nodes for this node.
117+
* @param dimensionValue : Value to match
118+
* @param lastMatchedChild : If not null, binary search will use this as the start/low
119+
* @return : Matched StarTreeNode or null if not found
120+
* @throws IOException : Any exception in reading the node data from index.
121+
*/
113122
StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException;
114123

115124
/**
116125
* Collects all matching child nodes whose dimension values lie within the range of low and high, both inclusive.
117126
* @param low : Starting of the range ( inclusive )
118127
* @param high : End of the range ( inclusive )
119128
* @param collector : Collector to collect the matched child StarTreeNode's
120-
* @throws IOException :
129+
* @throws IOException : Any exception in reading the node data from index.
121130
*/
122131
void collectChildrenInRange(long low, long high, StarTreeNodeCollector collector) throws IOException;
123132

‎server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ public boolean advanceExact(int target) throws IOException {
3939
return ((SortedSetDocValues) docIdSetIterator).advanceExact(target);
4040
}
4141

42+
// TODO : Remove this and merge @org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator to use value()
4243
public long nextOrd() throws IOException {
4344
return ((SortedSetDocValues) docIdSetIterator).nextOrd();
4445
}

‎server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java

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

99
package org.opensearch.search.startree;
1010

11-
import org.apache.lucene.index.LeafReaderContext;
1211
import org.apache.lucene.util.FixedBitSet;
1312
import org.opensearch.common.annotation.ExperimentalApi;
1413
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
@@ -63,6 +62,7 @@ public StarTreeQueryContext(SearchContext context, QueryBuilder baseQueryBuilder
6362
}
6463
}
6564

65+
// TODO : Make changes to change visibility into package private. Handle the same in @org.opensearch.search.SearchServiceStarTreeTests
6666
public StarTreeQueryContext(CompositeDataCubeFieldType compositeMappedFieldType, QueryBuilder baseQueryBuilder, int cacheSize) {
6767
this.compositeMappedFieldType = compositeMappedFieldType;
6868
this.baseQueryBuilder = baseQueryBuilder;
@@ -77,17 +77,17 @@ public CompositeIndexFieldInfo getStarTree() {
7777
return new CompositeIndexFieldInfo(compositeMappedFieldType.name(), compositeMappedFieldType.getCompositeIndexType());
7878
}
7979

80-
public FixedBitSet getStarTreeValue(LeafReaderContext ctx) {
81-
return starTreeValues != null ? starTreeValues[ctx.ord] : null;
80+
public FixedBitSet maybeGetCachedNodeIdsForSegment(int ordinal) {
81+
return starTreeValues != null ? starTreeValues[ordinal] : null;
8282
}
8383

84-
public FixedBitSet[] getStarTreeValues() {
84+
public FixedBitSet[] getAllCachedValues() {
8585
return starTreeValues;
8686
}
8787

88-
public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) {
88+
public void maybeSetCachedNodeIdsForSegment(int key, FixedBitSet values) {
8989
if (starTreeValues != null) {
90-
starTreeValues[ctx.ord] = values;
90+
starTreeValues[key] = values;
9191
}
9292
}
9393

‎server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,13 +80,11 @@ public static LeafBucketCollector getStarTreeLeafCollector(
8080
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
8181
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
8282

83-
assert starTreeValues != null;
8483
SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
8584
metricName
8685
);
8786
// Obtain a FixedBitSet of matched star tree document IDs
8887
FixedBitSet filteredValues = getStarTreeFilteredValues(context, ctx, starTreeValues);
89-
assert filteredValues != null;
9088

9189
int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs)
9290
if (numBits > 0) {
@@ -126,19 +124,19 @@ public void collect(int doc, long bucket) {
126124
*/
127125
public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
128126
throws IOException {
129-
FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().getStarTreeValue(ctx);
127+
FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().maybeGetCachedNodeIdsForSegment(ctx.ord);
130128
if (result == null) {
131129
result = StarTreeTraversalUtil.getStarTreeResult(
132130
starTreeValues,
133131
context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
134132
context
135133
);
136134
}
137-
context.getQueryShardContext().getStarTreeQueryContext().setStarTreeValues(ctx, result);
135+
context.getQueryShardContext().getStarTreeQueryContext().maybeSetCachedNodeIdsForSegment(ctx.ord, result);
138136
return result;
139137
}
140138

141-
public static Dimension getMatchingDimensionOrError(String dimensionName, List<Dimension> orderedDimensions) {
139+
public static Dimension getMatchingDimensionOrThrow(String dimensionName, List<Dimension> orderedDimensions) {
142140
Dimension matchingDimension = getMatchingDimensionOrNull(dimensionName, orderedDimensions);
143141
if (matchingDimension == null) {
144142
throw new IllegalStateException("No matching dimension found for [" + dimensionName + "]");

‎server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
4343
@Override
4444
public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) {
4545
convertedOrdinals = new TreeSet<>();
46-
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(
46+
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrThrow(
4747
dimensionName,
4848
starTreeValues.getStarTreeField().getDimensionsOrder()
4949
);

‎server/src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java

Lines changed: 33 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@
88

99
package org.opensearch.search.startree.filter.provider;
1010

11-
import org.apache.lucene.search.MatchNoDocsQuery;
12-
import org.apache.lucene.search.Query;
1311
import org.opensearch.common.annotation.ExperimentalApi;
1412
import org.opensearch.index.compositeindex.datacube.Dimension;
1513
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
@@ -53,20 +51,20 @@ StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, Composit
5351
*/
5452
class SingletonFactory {
5553

56-
private static final Map<Class<? extends QueryBuilder>, StarTreeFilterProvider> QUERY_BUILDERS_TO_STF_PROVIDER = Map.of(
57-
MatchAllQueryBuilder.class,
54+
private static final Map<String, StarTreeFilterProvider> QUERY_BUILDERS_TO_STF_PROVIDER = Map.of(
55+
MatchAllQueryBuilder.NAME,
5856
MATCH_ALL_PROVIDER,
59-
TermQueryBuilder.class,
57+
TermQueryBuilder.NAME,
6058
new TermStarTreeFilterProvider(),
61-
TermsQueryBuilder.class,
59+
TermsQueryBuilder.NAME,
6260
new TermsStarTreeFilterProvider(),
63-
RangeQueryBuilder.class,
61+
RangeQueryBuilder.NAME,
6462
new RangeStarTreeFilterProvider()
6563
);
6664

6765
public static StarTreeFilterProvider getProvider(QueryBuilder query) {
6866
if (query != null) {
69-
return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getClass());
67+
return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getName());
7068
}
7169
return MATCH_ALL_PROVIDER;
7270
}
@@ -83,23 +81,16 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
8381
TermQueryBuilder termQueryBuilder = (TermQueryBuilder) rawFilter;
8482
String field = termQueryBuilder.fieldName();
8583
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
86-
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
84+
DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null
85+
? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType)
86+
: null;
8787
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
8888
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
8989
return null; // Indicates Aggregators to fallback to default implementation.
9090
} else {
91-
// FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere.
92-
Query query = termQueryBuilder.toQuery(context.getQueryShardContext());
93-
if (query instanceof MatchNoDocsQuery) {
94-
return new StarTreeFilter(Collections.emptyMap());
95-
} else {
96-
return new StarTreeFilter(
97-
Map.of(
98-
field,
99-
List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value())))
100-
)
101-
);
102-
}
91+
return new StarTreeFilter(
92+
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value()))))
93+
);
10394
}
10495
}
10596
}
@@ -115,18 +106,15 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
115106
String field = termsQueryBuilder.fieldName();
116107
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
117108
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
118-
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
109+
DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null
110+
? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType)
111+
: null;
119112
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
120113
return null; // Indicates Aggregators to fallback to default implementation.
121114
} else {
122-
Query query = termsQueryBuilder.toQuery(context.getQueryShardContext());
123-
if (query instanceof MatchNoDocsQuery) {
124-
return new StarTreeFilter(Collections.emptyMap());
125-
} else {
126-
return new StarTreeFilter(
127-
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values())))
128-
);
129-
}
115+
return new StarTreeFilter(
116+
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values())))
117+
);
130118
}
131119
}
132120
}
@@ -143,29 +131,26 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
143131
String field = rangeQueryBuilder.fieldName();
144132
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
145133
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
146-
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
134+
DimensionFilterMapper dimensionFilterMapper = mappedFieldType == null
135+
? null
136+
: DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
147137
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
148138
return null;
149139
} else {
150-
Query query = rangeQueryBuilder.toQuery(context.getQueryShardContext());
151-
if (query instanceof MatchNoDocsQuery) {
152-
return new StarTreeFilter(Collections.emptyMap());
153-
} else {
154-
return new StarTreeFilter(
155-
Map.of(
156-
field,
157-
List.of(
158-
dimensionFilterMapper.getRangeMatchFilter(
159-
mappedFieldType,
160-
rangeQueryBuilder.from(),
161-
rangeQueryBuilder.to(),
162-
rangeQueryBuilder.includeLower(),
163-
rangeQueryBuilder.includeUpper()
164-
)
140+
return new StarTreeFilter(
141+
Map.of(
142+
field,
143+
List.of(
144+
dimensionFilterMapper.getRangeMatchFilter(
145+
mappedFieldType,
146+
rangeQueryBuilder.from(),
147+
rangeQueryBuilder.to(),
148+
rangeQueryBuilder.includeLower(),
149+
rangeQueryBuilder.includeUpper()
165150
)
166151
)
167-
);
168-
}
152+
)
153+
);
169154
}
170155
}
171156

‎server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeSearchTests.java

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,10 @@
2424
import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata;
2525
import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode;
2626
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeFactory;
27-
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
28-
import org.opensearch.search.startree.StarTreeNodeCollector;
27+
import org.opensearch.search.aggregations.startree.ArrayBasedCollector;
2928
import org.opensearch.test.OpenSearchTestCase;
3029

3130
import java.io.IOException;
32-
import java.util.ArrayList;
3331
import java.util.Arrays;
3432
import java.util.List;
3533
import java.util.TreeSet;
@@ -314,27 +312,27 @@ private void createStarTreeForDimension(
314312

315313
}
316314

317-
private static class ArrayBasedCollector implements StarTreeNodeCollector {
318-
319-
private final List<StarTreeNode> nodes = new ArrayList<>();
320-
321-
@Override
322-
public void collectStarTreeNode(StarTreeNode node) {
323-
nodes.add(node);
324-
}
325-
326-
public boolean matchValues(long[] values) throws IOException {
327-
boolean matches = true;
328-
for (int i = 0; i < values.length; i++) {
329-
matches &= nodes.get(i).getDimensionValue() == values[i];
330-
}
331-
return matches;
332-
}
333-
334-
public int collectedNodeCount() {
335-
return nodes.size();
336-
}
337-
338-
}
315+
// private static class ArrayBasedCollector implements StarTreeNodeCollector {
316+
//
317+
// private final List<StarTreeNode> nodes = new ArrayList<>();
318+
//
319+
// @Override
320+
// public void collectStarTreeNode(StarTreeNode node) {
321+
// nodes.add(node);
322+
// }
323+
//
324+
// public boolean matchValues(long[] values) throws IOException {
325+
// boolean matches = true;
326+
// for (int i = 0; i < values.length; i++) {
327+
// matches &= nodes.get(i).getDimensionValue() == values[i];
328+
// }
329+
// return matches;
330+
// }
331+
//
332+
// public int collectedNodeCount() {
333+
// return nodes.size();
334+
// }
335+
//
336+
// }
339337

340338
}

0 commit comments

Comments
 (0)