[Star Tree] [Search] Support of Date Range Queries in Aggregations

CHANGELOG.md

@@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added File Cache Pinning ([#17617](https://github.com/opensearch-project/OpenSearch/issues/13648))
 - Support consumer reset in Resume API for pull-based ingestion. This PR includes a breaking change for the experimental pull-based ingestion feature. ([#18332](https://github.com/opensearch-project/OpenSearch/pull/18332))
 - Add FIPS build tooling ([#4254](https://github.com/opensearch-project/security/issues/4254))
+- [Star-Tree] Support for date-range queries with star-tree supported aggregations ([#17855](https://github.com/opensearch-project/OpenSearch/pull/17855)
 
 ### Changed
 - Create generic DocRequest to better categorize ActionRequests ([#18269](https://github.com/opensearch-project/OpenSearch/pull/18269)))

server/src/main/java/org/opensearch/index/query/RangeQueryBuilder.java

@@ -35,6 +35,7 @@
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.geo.ShapeRelation;
 import org.opensearch.common.time.DateFormatter;
 import org.opensearch.common.time.DateMathParser;
@@ -56,10 +57,12 @@
 import java.util.Objects;
 
 /**
- * A Query that matches documents within an range of terms.
+ * A Query that matches documents within a range of terms.
  *
  * @opensearch.internal
  */
+// TODO: Revert @PublicApi annotation after @ExperimentalApi is removed from DimensionFilterMapper
+@PublicApi(since = "1.0.0")
 public class RangeQueryBuilder extends AbstractQueryBuilder<RangeQueryBuilder> implements MultiTermQueryBuilder {
     public static final String NAME = "range";
 

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -66,7 +66,7 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
 import org.opensearch.search.startree.StarTreeTraversalUtil;
-import org.opensearch.search.startree.filter.DimensionFilter;
+import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -301,8 +301,8 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
                 starTreeValues,
                 StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
                     context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
-                    starTreeDateDimension,
-                    List.of(DimensionFilter.MATCH_ALL_DEFAULT)
+                    fieldName,
+                    List.of(new MatchAllFilter(fieldName, starTreeDateDimension))
                 ),
                 context
             )

server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -71,7 +71,7 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
 import org.opensearch.search.startree.StarTreeTraversalUtil;
-import org.opensearch.search.startree.filter.DimensionFilter;
+import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -396,7 +396,7 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
                 StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
                     context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
                     fieldName,
-                    List.of(DimensionFilter.MATCH_ALL_DEFAULT)
+                    List.of(new MatchAllFilter(fieldName))
                 ),
                 context
             )

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -76,7 +76,7 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
 import org.opensearch.search.startree.StarTreeTraversalUtil;
-import org.opensearch.search.startree.filter.DimensionFilter;
+import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -343,7 +343,7 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
                 StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
                     context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
                     fieldName,
-                    List.of(DimensionFilter.MATCH_ALL_DEFAULT)
+                    List.of(new MatchAllFilter(fieldName))
                 ),
                 context
             )

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java

@@ -68,7 +68,7 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
 import org.opensearch.search.startree.StarTreeTraversalUtil;
-import org.opensearch.search.startree.filter.DimensionFilter;
+import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
 import java.math.BigInteger;
@@ -187,7 +187,7 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
                 StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
                     context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
                     fieldName,
-                    List.of(DimensionFilter.MATCH_ALL_DEFAULT)
+                    List.of(new MatchAllFilter(fieldName))
                 ),
                 context
             )

server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java

@@ -84,9 +84,11 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, StarT
         for (String remainingPredicateColumn : starTreeResult.remainingPredicateColumns) {
             logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
 
-            StarTreeValuesIterator valuesIterator = starTreeValues.getDimensionValuesIterator(remainingPredicateColumn);
             // Get the query value directly
             List<DimensionFilter> dimensionFilters = starTreeFilter.getFiltersForDimension(remainingPredicateColumn);
+            StarTreeValuesIterator valuesIterator = starTreeValues.getDimensionValuesIterator(
+                dimensionFilters.getFirst().getMatchingDimension()
+            );
 
             // Clear the temporary bit set before reuse
             tempBitSet.clear(0, starTreeResult.maxMatchedDoc + 1);
@@ -131,7 +133,7 @@ private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, St
         Queue<StarTreeNode> queue = new ArrayDeque<>();
         queue.add(starTree);
         int currentDimensionId = -1;
-        Set<String> remainingPredicateColumns = new HashSet<>(starTreeFilter.getDimensions());
+        Set<String> remainingPredicateColumns = new HashSet<>(starTreeFilter.getMatchingDimensions());
         int matchedDocsCountInStarTree = 0;
         int maxDocNum = -1;
         StarTreeNode starTreeNode;

server/src/main/java/org/opensearch/search/startree/filter/DimensionFilter.java

@@ -11,44 +11,16 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
-import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeNodeCollector;
 
 import java.io.IOException;
-import java.util.Iterator;
 
 /**
  * Contains the logic to filter over a dimension either in StarTree Index or it's Dimension DocValues
  */
 @ExperimentalApi
 public interface DimensionFilter {
-
-    DimensionFilter MATCH_ALL_DEFAULT = new DimensionFilter() {
-        @Override
-        public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) throws IOException {
-
-        }
-
-        @Override
-        public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector)
-            throws IOException {
-            if (parentNode != null) {
-                for (Iterator<? extends StarTreeNode> it = parentNode.getChildrenIterator(); it.hasNext();) {
-                    StarTreeNode starTreeNode = it.next();
-                    if (starTreeNode.getStarTreeNodeType() == StarTreeNodeType.DEFAULT.getValue()) {
-                        collector.collectStarTreeNode(starTreeNode);
-                    }
-                }
-            }
-        }
-
-        @Override
-        public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
-            return true;
-        }
-    };
-
     /**
      * Converts parsed user values to ordinals based on segment and other init actions can be performed.
      * @param starTreeValues : Segment specific star tree root node and other metadata
@@ -74,10 +46,16 @@ public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
      */
     boolean matchDimValue(long ordinal, StarTreeValues starTreeValues);
 
-    default String getDimensionName() {
+    String getDimensionName();
+
+    default String getSubDimensionName() {
         return null;
     }
 
+    default String getMatchingDimension() {
+        return getSubDimensionName() == null ? getDimensionName() : getSubDimensionName();
+    }
+
     /**
      * Represents how to match a value when comparing during StarTreeTraversal
      */
@@ -89,5 +67,4 @@ enum MatchType {
         LTE,
         EXACT
     }
-
 }

server/src/main/java/org/opensearch/search/startree/filter/DimensionFilterMergerUtils.java

@@ -133,7 +133,16 @@ private static DimensionFilter intersectRangeFilters(
             }
         }
 
-        return new RangeMatchDimFilter(range1.getDimensionName(), newLow, newHigh, includeLow, includeHigh);
+        String effectiveSubDimension = mapper.resolveUsingSubDimension()
+            ? mapper.getSubDimensionFieldEffective(range1.getSubDimensionName(), range2.getSubDimensionName())
+            : null;
+
+        return new RangeMatchDimFilter(range1.getDimensionName(), newLow, newHigh, includeLow, includeHigh) {
+            @Override
+            public String getSubDimensionName() {
+                return effectiveSubDimension;
+            }
+        };
     }
 
     /**

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

@@ -50,7 +50,8 @@ public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext se
             starTreeValues.getStarTreeField().getDimensionsOrder()
         );
         this.dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
-            searchContext.mapperService().fieldType(dimensionName)
+            searchContext.mapperService().fieldType(dimensionName),
+            searchContext
         );
 
         for (Object rawValue : rawValues) {

server/src/main/java/org/opensearch/search/startree/filter/MatchAllFilter.java

@@ -0,0 +1,70 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.startree.filter;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
+import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.StarTreeNodeCollector;
+
+import java.io.IOException;
+import java.util.Iterator;
+
+/**
+ * Matches all StarTreeNodes
+ */
+@ExperimentalApi
+public class MatchAllFilter implements DimensionFilter {
+
+    public final String dimensionName;
+    public final String subDimensionName;
+
+    public MatchAllFilter(String dimensionName) {
+        this.dimensionName = dimensionName;
+        this.subDimensionName = null;
+    }
+
+    public MatchAllFilter(String dimensionName, String subDimensionName) {
+        this.dimensionName = dimensionName;
+        this.subDimensionName = subDimensionName;
+    }
+
+    @Override
+    public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) throws IOException {}
+
+    @Override
+    public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeValues, StarTreeNodeCollector collector)
+        throws IOException {
+        if (parentNode != null) {
+            for (Iterator<? extends StarTreeNode> it = parentNode.getChildrenIterator(); it.hasNext();) {
+                StarTreeNode starTreeNode = it.next();
+                if (starTreeNode.getStarTreeNodeType() == StarTreeNodeType.DEFAULT.getValue()) {
+                    collector.collectStarTreeNode(starTreeNode);
+                }
+            }
+        }
+    }
+
+    @Override
+    public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
+        return true;
+    }
+
+    @Override
+    public String getDimensionName() {
+        return dimensionName;
+    }
+
+    public String getSubDimensionName() {
+        return subDimensionName;
+    }
+
+}

server/src/main/java/org/opensearch/search/startree/filter/MatchNoneFilter.java

@@ -34,4 +34,9 @@
     public boolean matchDimValue(long ordinal, StarTreeValues starTreeValues) {
         return false;
     }
+
+    @Override
+    public String getDimensionName() {
+        return null;
+    }
 }

server/src/main/java/org/opensearch/search/startree/filter/RangeMatchDimFilter.java

@@ -51,7 +51,8 @@ public RangeMatchDimFilter(String dimensionName, Object low, Object high, boolea
     public void initialiseForSegment(StarTreeValues starTreeValues, SearchContext searchContext) {
         skipRangeCollection = false;
         this.dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
-            searchContext.mapperService().fieldType(dimensionName)
+            searchContext.mapperService().fieldType(dimensionName),
+            searchContext
         );
 
         lowOrdinal = 0L;

server/src/main/java/org/opensearch/search/startree/filter/StarTreeFilter.java

@@ -10,6 +10,7 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -23,16 +24,38 @@ public class StarTreeFilter {
     private final Map<String, List<DimensionFilter>> dimensionFilterMap;
 
     public StarTreeFilter(Map<String, List<DimensionFilter>> dimensionFilterMap) {
+        // TODO: Evaluate wrapping of List<DimensionFilter> to expose a single effective sub-dimension field
+        // https://github.com/opensearch-project/OpenSearch/issues/18477
         this.dimensionFilterMap = dimensionFilterMap;
     }
 
     public List<DimensionFilter> getFiltersForDimension(String dimension) {
-        return dimensionFilterMap.get(dimension);
+        // TODO: Refactor with https://github.com/opensearch-project/OpenSearch/issues/18477
+        return dimensionFilterMap.containsKey(dimension)
+            ? dimensionFilterMap.get(dimension)
+            : dimensionFilterMap.values()
+                .stream()
+                .filter(filters -> filters != null && !filters.isEmpty() && dimension.equals(filters.getFirst().getMatchingDimension()))
+                .findFirst()
+                .orElse(null);
     }
 
     public Set<String> getDimensions() {
         return dimensionFilterMap.keySet();
     }
+
+    public Set<String> getMatchingDimensions() {
+        // TODO: Refactor with: https://github.com/opensearch-project/OpenSearch/issues/18477
+        Set<String> matchingDimensions = new HashSet<>();
+        for (Map.Entry<String, List<DimensionFilter>> entry : dimensionFilterMap.entrySet()) {
+            if (entry.getValue() != null && !entry.getValue().isEmpty()) {
+                String matchingDimension = entry.getValue().getFirst().getMatchingDimension();
+                // Anonymous implementations of DimensionFilter may not have dimension info
+                matchingDimensions.add(matchingDimension == null ? entry.getKey() : matchingDimension);
+            }
+        }
+        return matchingDimensions;
+    }
     // TODO : Implement Merging of 2 Star Tree Filters
     // This would also involve merging 2 different types of dimension filters.
     // It also brings in the challenge of sorting input values in user query for efficient merging.