Skip to content

Commit

Permalink
Increased UT coverage and addressed PR Comments
Browse files Browse the repository at this point in the history 
Signed-off-by: expani <anijainc@amazon.com>
  • Loading branch information
@expani
expani committed Jan 28, 2025
1 parent 8d1204f commit e386545
Show file tree
Hide file tree
Showing 11 changed files with 267 additions and 106 deletions.
14 changes: 0 additions & 14 deletions ...arch/index/compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNode.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -192,20 +192,6 @@ public StarTreeNode getChildStarNode() throws IOException {
return handleStarNode();
}

@Override
public StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
// there will be no children for leaf nodes
if (isLeaf()) {
return null;
}

StarTreeNode resultStarTreeNode = null;
if (null != dimensionValue) {
resultStarTreeNode = binarySearchChild(dimensionValue, null);
}
return resultStarTreeNode;
}

@Override
public StarTreeNode getChildForDimensionValue(Long dimensionValue, StarTreeNode lastMatchedChild) throws IOException {
// there will be no children for leaf nodes
Expand Down
13 changes: 11 additions & 2 deletions ...rc/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,16 +108,25 @@ public interface StarTreeNode {
* @return the child node for the given dimension value or null if child is not present
* @throws IOException if an I/O error occurs while retrieving the child node
*/
StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException;
default StarTreeNode getChildForDimensionValue(Long dimensionValue) throws IOException {
return getChildForDimensionValue(dimensionValue, null);
}

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

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

Expand Down
1 change: 1 addition & 0 deletions ...ndex/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ public boolean advanceExact(int target) throws IOException {
return ((SortedSetDocValues) docIdSetIterator).advanceExact(target);
}

// TODO : Remove this and merge @org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator to use value()
public long nextOrd() throws IOException {
return ((SortedSetDocValues) docIdSetIterator).nextOrd();
}
Expand Down
12 changes: 6 additions & 6 deletions server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

package org.opensearch.search.startree;

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
Expand Down Expand Up @@ -68,6 +67,7 @@ public StarTreeQueryContext(SearchContext context, QueryBuilder baseQueryBuilder
}
}

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

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

public FixedBitSet[] getStarTreeValues() {
public FixedBitSet[] getAllCachedValues() {
return starTreeValues;
}

public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) {
public void maybeSetCachedNodeIdsForSegment(int key, FixedBitSet values) {
if (starTreeValues != null) {
starTreeValues[ctx.ord] = values;
starTreeValues[key] = values;
}
}

Expand Down
8 changes: 3 additions & 5 deletions server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,11 @@ public static LeafBucketCollector getStarTreeLeafCollector(
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);

assert starTreeValues != null;
SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
metricName
);
// Obtain a FixedBitSet of matched star tree document IDs
FixedBitSet filteredValues = getStarTreeFilteredValues(context, ctx, starTreeValues);
assert filteredValues != null;

int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs)
if (numBits > 0) {
Expand Down Expand Up @@ -122,19 +120,19 @@ public static LeafBucketCollector getStarTreeLeafCollector(
*/
public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
throws IOException {
FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().getStarTreeValue(ctx);
FixedBitSet result = context.getQueryShardContext().getStarTreeQueryContext().maybeGetCachedNodeIdsForSegment(ctx.ord);
if (result == null) {
result = StarTreeTraversalUtil.getStarTreeResult(
starTreeValues,
context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
context
);
}
context.getQueryShardContext().getStarTreeQueryContext().setStarTreeValues(ctx, result);
context.getQueryShardContext().getStarTreeQueryContext().maybeSetCachedNodeIdsForSegment(ctx.ord, result);
return result;
}

public static Dimension getMatchingDimensionOrError(String dimensionName, List<Dimension> orderedDimensions) {
public static Dimension getMatchingDimensionOrThrow(String dimensionName, List<Dimension> orderedDimensions) {
Dimension matchingDimension = getMatchingDimensionOrNull(dimensionName, orderedDimensions);
if (matchingDimension == null) {
throw new IllegalStateException("No matching dimension found for [" + dimensionName + "]");
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/search/startree/filter/ExactMatchDimFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
@Override
public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) {
convertedOrdinals = new TreeSet<>();
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrThrow(
dimensionName,
starTreeValues.getStarTreeField().getDimensionsOrder()
);
Expand Down
81 changes: 33 additions & 48 deletions .../src/main/java/org/opensearch/search/startree/filter/provider/StarTreeFilterProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,6 @@

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

import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
Expand Down Expand Up @@ -53,20 +51,20 @@ StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, Composit
*/
class SingletonFactory {

private static final Map<Class<? extends QueryBuilder>, StarTreeFilterProvider> QUERY_BUILDERS_TO_STF_PROVIDER = Map.of(
MatchAllQueryBuilder.class,
private static final Map<String, StarTreeFilterProvider> QUERY_BUILDERS_TO_STF_PROVIDER = Map.of(
MatchAllQueryBuilder.NAME,
MATCH_ALL_PROVIDER,
TermQueryBuilder.class,
TermQueryBuilder.NAME,
new TermStarTreeFilterProvider(),
TermsQueryBuilder.class,
TermsQueryBuilder.NAME,
new TermsStarTreeFilterProvider(),
RangeQueryBuilder.class,
RangeQueryBuilder.NAME,
new RangeStarTreeFilterProvider()
);

public static StarTreeFilterProvider getProvider(QueryBuilder query) {
if (query != null) {
return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getClass());
return QUERY_BUILDERS_TO_STF_PROVIDER.get(query.getName());
}
return MATCH_ALL_PROVIDER;
}
Expand All @@ -83,23 +81,16 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
TermQueryBuilder termQueryBuilder = (TermQueryBuilder) rawFilter;
String field = termQueryBuilder.fieldName();
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null
? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType)
: null;
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
return null; // Indicates Aggregators to fallback to default implementation.
} else {
// FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere.
Query query = termQueryBuilder.toQuery(context.getQueryShardContext());
if (query instanceof MatchNoDocsQuery) {
return new StarTreeFilter(Collections.emptyMap());
} else {
return new StarTreeFilter(
Map.of(
field,
List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value())))
)
);
}
return new StarTreeFilter(
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, List.of(termQueryBuilder.value()))))
);
}
}
}
Expand All @@ -115,18 +106,15 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
String field = termsQueryBuilder.fieldName();
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
DimensionFilterMapper dimensionFilterMapper = mappedFieldType != null
? DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType)
: null;
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
return null; // Indicates Aggregators to fallback to default implementation.
} else {
Query query = termsQueryBuilder.toQuery(context.getQueryShardContext());
if (query instanceof MatchNoDocsQuery) {
return new StarTreeFilter(Collections.emptyMap());
} else {
return new StarTreeFilter(
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values())))
);
}
return new StarTreeFilter(
Map.of(field, List.of(dimensionFilterMapper.getExactMatchFilter(mappedFieldType, termsQueryBuilder.values())))
);
}
}
}
Expand All @@ -143,29 +131,26 @@ public StarTreeFilter getFilter(SearchContext context, QueryBuilder rawFilter, C
String field = rangeQueryBuilder.fieldName();
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
MappedFieldType mappedFieldType = context.mapperService().fieldType(field);
DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
DimensionFilterMapper dimensionFilterMapper = mappedFieldType == null
? null
: DimensionFilterMapper.Factory.fromMappedFieldType(mappedFieldType);
if (matchedDimension == null || mappedFieldType == null || dimensionFilterMapper == null) {
return null;
} else {
Query query = rangeQueryBuilder.toQuery(context.getQueryShardContext());
if (query instanceof MatchNoDocsQuery) {
return new StarTreeFilter(Collections.emptyMap());
} else {
return new StarTreeFilter(
Map.of(
field,
List.of(
dimensionFilterMapper.getRangeMatchFilter(
mappedFieldType,
rangeQueryBuilder.from(),
rangeQueryBuilder.to(),
rangeQueryBuilder.includeLower(),
rangeQueryBuilder.includeUpper()
)
return new StarTreeFilter(
Map.of(
field,
List.of(
dimensionFilterMapper.getRangeMatchFilter(
mappedFieldType,
rangeQueryBuilder.from(),
rangeQueryBuilder.to(),
rangeQueryBuilder.includeLower(),
rangeQueryBuilder.includeUpper()
)
)
);
}
)
);
}
}

Expand Down
48 changes: 23 additions & 25 deletions ...compositeindex/datacube/startree/fileformats/node/FixedLengthStarTreeNodeSearchTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,10 @@
import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata;
import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeFactory;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
import org.opensearch.search.startree.StarTreeNodeCollector;
import org.opensearch.search.aggregations.startree.ArrayBasedCollector;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.TreeSet;
Expand Down Expand Up @@ -314,27 +312,27 @@ private void createStarTreeForDimension(

}

private static class ArrayBasedCollector implements StarTreeNodeCollector {

private final List<StarTreeNode> nodes = new ArrayList<>();

@Override
public void collectStarTreeNode(StarTreeNode node) {
nodes.add(node);
}

public boolean matchValues(long[] values) throws IOException {
boolean matches = true;
for (int i = 0; i < values.length; i++) {
matches &= nodes.get(i).getDimensionValue() == values[i];
}
return matches;
}

public int collectedNodeCount() {
return nodes.size();
}

}
// private static class ArrayBasedCollector implements StarTreeNodeCollector {
//
// private final List<StarTreeNode> nodes = new ArrayList<>();
//
// @Override
// public void collectStarTreeNode(StarTreeNode node) {
// nodes.add(node);
// }
//
// public boolean matchValues(long[] values) throws IOException {
// boolean matches = true;
// for (int i = 0; i < values.length; i++) {
// matches &= nodes.get(i).getDimensionValue() == values[i];
// }
// return matches;
// }
//
// public int collectedNodeCount() {
// return nodes.size();
// }
//
// }

}
Loading

0 comments on commit e386545

Please sign in to comment.