From b42f5e9fb3a2689ba81ed8479ce3e58b74453c29 Mon Sep 17 00:00:00 2001 From: expani Date: Thu, 16 Jan 2025 18:34:44 -0800 Subject: [PATCH] Refactoring getDimension and changed ordinal fetch for keyword Signed-off-by: expani --- .../search/startree/ExactMatchDimFilter.java | 5 +- .../FieldToDimensionOrdinalMapper.java | 6 +- .../search/startree/RangeMatchDimFilter.java | 35 +++++---- .../startree/StarTreeFilterProvider.java | 71 +++++++++---------- .../startree/StarTreeNodeCollector.java | 2 +- .../search/startree/StarTreeQueryHelper.java | 16 +++-- 6 files changed, 71 insertions(+), 64 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java b/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java index a3230de8bc8f1..523a6c17ccee5 100644 --- a/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/ExactMatchDimFilter.java @@ -36,7 +36,10 @@ public ExactMatchDimFilter(String dimensionName, List 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( diff --git a/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java b/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java index 0c725c1f464cd..c8b816d85459f 100644 --- a/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java +++ b/server/src/main/java/org/opensearch/search/startree/FieldToDimensionOrdinalMapper.java @@ -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) { @@ -101,7 +101,7 @@ enum MatchType { LT, GTE, LTE, - EXACT; + EXACT } } diff --git a/server/src/main/java/org/opensearch/search/startree/RangeMatchDimFilter.java b/server/src/main/java/org/opensearch/search/startree/RangeMatchDimFilter.java index 59586afd09f13..2243e7294ae79 100644 --- a/server/src/main/java/org/opensearch/search/startree/RangeMatchDimFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/RangeMatchDimFilter.java @@ -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; @@ -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; + } } } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilterProvider.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilterProvider.java index 17a4dabab2866..996e9c6321d89 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilterProvider.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilterProvider.java @@ -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; @@ -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 { @@ -33,57 +32,44 @@ class SingletonFactory { (rawFilter, compositeFieldType) -> { TermQueryBuilder termQueryBuilder = (TermQueryBuilder) rawFilter; String field = termQueryBuilder.fieldName(); - List 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 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 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() + ) ) ) - ) - ); + ); } ); @@ -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; + } + } } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeNodeCollector.java b/server/src/main/java/org/opensearch/search/startree/StarTreeNodeCollector.java index a11109dba4385..750403c50527a 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeNodeCollector.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeNodeCollector.java @@ -14,6 +14,6 @@ @ExperimentalApi public interface StarTreeNodeCollector { - public void collectStarNode(StarTreeNode node); + void collectStarNode(StarTreeNode node); } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java index 722287494ade9..8d3ec0aec6925 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java @@ -360,14 +360,20 @@ private static StarTreeNode reachClosestParent(StarTreeNode startNode, int dimen return currentNode; } - public static Dimension getMatchingDimensionOrError(String dimensionName, StarTreeValues starTreeValues) { - List matchingDimensions = starTreeValues.getStarTreeField() - .getDimensionsOrder() - .stream() + public static Dimension getMatchingDimensionOrError(String dimensionName, List 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 orderedDimensions) { + List 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); }