Skip to content

Commit

Permalink
Refactoring getDimension and changed ordinal fetch for keyword
Browse files Browse the repository at this point in the history 
Signed-off-by: expani <[email protected]>
  • Loading branch information
@expani
expani committed Jan 17, 2025
1 parent 7df7ff7 commit b42f5e9
Show file tree
Hide file tree
Showing 6 changed files with 71 additions and 64 deletions.
5 changes: 4 additions & 1 deletion server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,10 @@ public ExactMatchDimFilter(String dimensionName, List<Object> valuesToMatch) {
@Override
public void initialiseForSegment(StarTreeValues starTreeValues) {
convertedOrdinals = new TreeSet<>();
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(dimensionName, starTreeValues);
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(
dimensionName,
starTreeValues.getStarTreeField().getDimensionsOrder()
);
FieldToDimensionOrdinalMapper fieldToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(matchedDim.getDocValuesType());
for (Object rawValue : rawValues) {
long ordinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(
Expand Down
6 changes: 3 additions & 3 deletions server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@ enum SingletonFactory {
TermsEnum.SeekStatus seekStatus = termsEnum.seekCeil((BytesRef) value);
if (matchType == MatchType.GT || matchType == MatchType.GTE) {
return termsEnum.ord();
} else {
return (seekStatus == TermsEnum.SeekStatus.FOUND) ? termsEnum.ord() : termsEnum.ord() - 1;
} else { // LT || LTE
return (seekStatus == TermsEnum.SeekStatus.NOT_FOUND) ? termsEnum.ord() - 1 : termsEnum.ord();
}
}
} catch (IOException e) {
Expand Down Expand Up @@ -101,7 +101,7 @@ enum MatchType {
LT,
GTE,
LTE,
EXACT;
EXACT
}

}
35 changes: 20 additions & 15 deletions server/src/main/java/org/opensearch/search/startree/RangeMatchDimFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ public class RangeMatchDimFilter implements DimensionFilter {
private final boolean includeLow;
private final boolean includeHigh;

private long lowOrdinal;
private long highOrdinal;
private Long lowOrdinal;
private Long highOrdinal;

public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolean includeLow, boolean includeHigh) {
this.dimensionName = dimensionName;
Expand All @@ -39,19 +39,24 @@ public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolea

@Override
public void initialiseForSegment(StarTreeValues starTreeValues) throws IOException {
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(dimensionName, starTreeValues);
FieldToDimensionOrdinalMapper fieldToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(matchedDim.getDocValuesType());
if (low != null) {
MatchType lowMatchType = includeLow ? MatchType.GTE : MatchType.GT;
lowOrdinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(dimensionName, low, starTreeValues, lowMatchType);
} else {
lowOrdinal = 0;
}
if (high != null) {
MatchType highMatchType = includeHigh ? MatchType.LTE : MatchType.LT;
highOrdinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(dimensionName, high, starTreeValues, highMatchType);
} else {
highOrdinal = Long.MAX_VALUE;
Dimension matchedDim = StarTreeQueryHelper.getMatchingDimensionOrError(
dimensionName,
starTreeValues.getStarTreeField().getDimensionsOrder()
);
if (lowOrdinal != null && highOrdinal != null) {
FieldToDimensionOrdinalMapper fieldToDimensionOrdinalMapper = getFieldToDimensionOrdinalMapper(matchedDim.getDocValuesType());
if (low != null) {
MatchType lowMatchType = includeLow ? MatchType.GTE : MatchType.GT;
lowOrdinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(dimensionName, low, starTreeValues, lowMatchType);
} else {
lowOrdinal = 0L;
}
if (high != null) {
MatchType highMatchType = includeHigh ? MatchType.LTE : MatchType.LT;
highOrdinal = fieldToDimensionOrdinalMapper.getMatchingOrdinal(dimensionName, high, starTreeValues, highMatchType);
} else {
highOrdinal = Long.MAX_VALUE;
}
}
}

Expand Down
71 changes: 32 additions & 39 deletions server/src/main/java/org/opensearch/search/startree/StarTreeFilterProvider.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.DocValuesType;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
Expand All @@ -17,14 +16,14 @@
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.index.query.TermsQueryBuilder;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@ExperimentalApi
public interface StarTreeFilterProvider {

public StarTreeFilter getFilter(QueryBuilder rawFilter, CompositeDataCubeFieldType compositeFieldType);
StarTreeFilter getFilter(QueryBuilder rawFilter, CompositeDataCubeFieldType compositeFieldType);

class SingletonFactory {

Expand All @@ -33,57 +32,44 @@ class SingletonFactory {
(rawFilter, compositeFieldType) -> {
TermQueryBuilder termQueryBuilder = (TermQueryBuilder) rawFilter;
String field = termQueryBuilder.fieldName();
List<Dimension> matchedDimension = compositeFieldType.getDimensions()
.stream()
.filter(dim -> dim.getField().equals(field))
.collect(Collectors.toList());
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
// FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere.
if (matchedDimension.size() != 1 || matchedDimension.get(0).getDocValuesType() != DocValuesType.SORTED_NUMERIC) {
return null;
}
return new StarTreeFilter(Map.of(field, List.of(new ExactMatchDimFilter(field, List.of(termQueryBuilder.value())))));
return matchedDimension == null
? new StarTreeFilter(Collections.emptyMap())
: new StarTreeFilter(Map.of(field, List.of(new ExactMatchDimFilter(field, List.of(termQueryBuilder.value())))));
},
TermsQueryBuilder.class,
(rawFilter, compositeFieldType) -> {
TermsQueryBuilder termsQueryBuilder = (TermsQueryBuilder) rawFilter;
String field = termsQueryBuilder.fieldName();
List<Dimension> matchedDimension = compositeFieldType.getDimensions()
.stream()
.filter(dim -> dim.getField().equals(field))
.collect(Collectors.toList());
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
// FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere.
if (matchedDimension.size() != 1 || matchedDimension.get(0).getDocValuesType() != DocValuesType.SORTED_NUMERIC) {
return null;
}
return new StarTreeFilter(Map.of(field, List.of(new ExactMatchDimFilter(field, termsQueryBuilder.values()))));
return matchedDimension == null
? new StarTreeFilter(Collections.emptyMap())
: new StarTreeFilter(Map.of(field, List.of(new ExactMatchDimFilter(field, termsQueryBuilder.values()))));
},
RangeQueryBuilder.class,
(rawFilter, compositeFieldType) -> {
RangeQueryBuilder rangeQueryBuilder = (RangeQueryBuilder) rawFilter;
String field = rangeQueryBuilder.fieldName();
List<Dimension> matchedDimensions = compositeFieldType.getDimensions()
.stream()
.filter(dim -> dim.getField().equals(field))
.collect(Collectors.toList());
Dimension matchedDimension = StarTreeQueryHelper.getMatchingDimensionOrNull(field, compositeFieldType.getDimensions());
// FIXME : DocValuesType validation is field type specific and not query builder specific should happen elsewhere.
if (matchedDimensions.size() != 1 || matchedDimensions.get(0).getDocValuesType() != DocValuesType.SORTED_NUMERIC) {
return null;
}
Dimension matchedDimension = matchedDimensions.get(0);
return new StarTreeFilter(
Map.of(
field,
List.of(
new RangeMatchDimFilter(
matchedDimension.getField(),
rangeQueryBuilder.from(),
rangeQueryBuilder.to(),
rangeQueryBuilder.includeLower(),
rangeQueryBuilder.includeUpper()
return matchedDimension == null
? new StarTreeFilter(Collections.emptyMap())
: new StarTreeFilter(
Map.of(
field,
List.of(
new RangeMatchDimFilter(
matchedDimension.getField(),
rangeQueryBuilder.from(),
rangeQueryBuilder.to(),
rangeQueryBuilder.includeLower(),
rangeQueryBuilder.includeUpper()
)
)
)
)
);
);
}
);

Expand All @@ -95,6 +81,13 @@ public static StarTreeFilterProvider getProvider(QueryBuilder queryBuilder) {
}
}

private static Dimension getDimensionOrError(Dimension dimension) {
if (dimension == null) {
throw new IllegalStateException("Field [" + dimension.getField() + "] not found in star tree");
}
return dimension;
}

}

}
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/search/startree/StarTreeNodeCollector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,6 @@
@ExperimentalApi
public interface StarTreeNodeCollector {

public void collectStarNode(StarTreeNode node);
void collectStarNode(StarTreeNode node);

}
16 changes: 11 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 @@ -360,14 +360,20 @@ private static StarTreeNode reachClosestParent(StarTreeNode startNode, int dimen
return currentNode;
}

public static Dimension getMatchingDimensionOrError(String dimensionName, StarTreeValues starTreeValues) {
List<Dimension> matchingDimensions = starTreeValues.getStarTreeField()
.getDimensionsOrder()
.stream()
public static Dimension getMatchingDimensionOrError(String dimensionName, List<Dimension> orderedDimensions) {
Dimension matchingDimension = getMatchingDimensionOrNull(dimensionName, orderedDimensions);
if (matchingDimension == null) {
throw new IllegalStateException("No matching dimension found for [" + dimensionName + "]");
}
return matchingDimension;
}

public static Dimension getMatchingDimensionOrNull(String dimensionName, List<Dimension> orderedDimensions) {
List<Dimension> matchingDimensions = orderedDimensions.stream()
.filter(x -> x.getField().equals(dimensionName))
.collect(Collectors.toList());
if (matchingDimensions.size() != 1) {
throw new IllegalStateException("Expected exactly one dimension but found " + matchingDimensions);
return null;
}
return matchingDimensions.get(0);
}
Expand Down

0 comments on commit b42f5e9

Please sign in to comment.