From a131fecbea9b199b8504138f9bda75c3d3558c7a Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Sun, 25 Aug 2024 02:59:43 -0700
Subject: [PATCH 01/35] Star Tree Search request/response changes
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeUtils.java | 24 ++
.../index/query/QueryShardContext.java | 77 +++++
.../org/opensearch/search/SearchService.java | 69 ++++-
.../aggregations/AggregatorFactories.java | 6 +-
.../aggregations/AggregatorFactory.java | 4 +
.../aggregations/metrics/AvgAggregator.java | 66 ++++
.../metrics/AvgAggregatorFactory.java | 2 +-
.../aggregations/metrics/MaxAggregator.java | 44 +++
.../metrics/MaxAggregatorFactory.java | 2 +-
.../aggregations/metrics/MinAggregator.java | 42 +++
.../metrics/MinAggregatorFactory.java | 2 +-
.../metrics/NumericMetricsAggregator.java | 19 ++
.../aggregations/metrics/SumAggregator.java | 59 +++-
.../metrics/SumAggregatorFactory.java | 2 +-
.../metrics/ValueCountAggregator.java | 34 +++
.../metrics/ValueCountAggregatorFactory.java | 2 +-
.../aggregations/support/ValuesSource.java | 4 +
.../ValuesSourceAggregatorFactory.java | 12 +
.../startree/OriginalOrStarTreeQuery.java | 83 +++++
.../search/startree/StarTreeFilter.java | 289 ++++++++++++++++++
.../search/startree/StarTreeQuery.java | 116 +++++++
.../search/startree/package-info.java | 10 +
.../StarTreeDocValuesFormatTests.java | 10 +-
.../startree/MetricAggregatorTests.java | 212 +++++++++++++
.../aggregations/AggregatorTestCase.java | 70 +++++
25 files changed, 1240 insertions(+), 20 deletions(-)
create mode 100644 server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
create mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
create mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
create mode 100644 server/src/main/java/org/opensearch/search/startree/package-info.java
create mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
index 2aae0d4ca7e29..201bb895ee49a 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
@@ -12,10 +12,18 @@
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.VectorEncoding;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo;
+import org.opensearch.search.aggregations.metrics.AvgAggregatorFactory;
+import org.opensearch.search.aggregations.metrics.MaxAggregatorFactory;
+import org.opensearch.search.aggregations.metrics.MinAggregatorFactory;
+import org.opensearch.search.aggregations.metrics.SumAggregatorFactory;
+import org.opensearch.search.aggregations.metrics.ValueCountAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import java.util.Collections;
import java.util.List;
+import java.util.Map;
/**
* Util class for building star tree
@@ -38,6 +46,22 @@ private StarTreeUtils() {}
*/
public static final String METRIC_SUFFIX = "metric";
+ /**
+ * Map to associate star-tree supported AggregatorFactory classes with their corresponding MetricStat
+ */
+ public static final Map<Class<? extends ValuesSourceAggregatorFactory>, MetricStat> aggregatorStatMap = Map.of(
+ SumAggregatorFactory.class,
+ MetricStat.SUM,
+ MaxAggregatorFactory.class,
+ MetricStat.MAX,
+ MinAggregatorFactory.class,
+ MetricStat.MIN,
+ ValueCountAggregatorFactory.class,
+ MetricStat.VALUE_COUNT,
+ AvgAggregatorFactory.class,
+ MetricStat.AVG
+ );
+
/**
* Returns the full field name for a dimension in the star-tree index.
*
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index bccead2b029d0..894f467a8f6fc 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -56,7 +56,13 @@
import org.opensearch.index.IndexSortConfig;
import org.opensearch.index.analysis.IndexAnalyzers;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.Dimension;
+import org.opensearch.index.compositeindex.datacube.Metric;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.IndexFieldData;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.ContentPath;
import org.opensearch.index.mapper.DerivedFieldResolver;
import org.opensearch.index.mapper.DerivedFieldResolverFactory;
@@ -73,12 +79,17 @@
import org.opensearch.script.ScriptContext;
import org.opensearch.script.ScriptFactory;
import org.opensearch.script.ScriptService;
+import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.support.AggregationUsageService;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.lookup.SearchLookup;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.transport.RemoteClusterAware;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
@@ -89,6 +100,7 @@
import java.util.function.LongSupplier;
import java.util.function.Predicate;
import java.util.function.Supplier;
+import java.util.stream.Collectors;
import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
@@ -582,6 +594,71 @@ private ParsedQuery toQuery(QueryBuilder queryBuilder, CheckedFunction<QueryBuil
}
}
+ public ParsedQuery toStarTreeQuery(
+ CompositeIndexFieldInfo starTree,
+ CompositeDataCubeFieldType compositeIndexFieldInfo,
+ QueryBuilder queryBuilder,
+ Query query
+ ) {
+ Map<String, List<Predicate<Long>>> predicateMap;
+
+ if (queryBuilder == null) {
+ predicateMap = null;
+ } else if (queryBuilder instanceof TermQueryBuilder) {
+ List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
+ .stream()
+ .map(Dimension::getField)
+ .collect(Collectors.toList());
+ predicateMap = getStarTreePredicates(queryBuilder, supportedDimensions);
+ } else {
+ return null;
+ }
+
+ StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, predicateMap);
+ OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, query);
+ return new ParsedQuery(originalOrStarTreeQuery);
+ }
+
+ /**
+ * Parse query body to star-tree predicates
+ * @param queryBuilder
+ * @return predicates to match
+ */
+ private Map<String, List<Predicate<Long>>> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
+ TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
+ String field = tq.fieldName();
+ if (supportedDimensions.contains(field) == false) {
+ throw new IllegalArgumentException("unsupported field in star-tree");
+ }
+ long inputQueryVal = Long.parseLong(tq.value().toString());
+
+ // Get or create the list of predicates for the given field
+ Map<String, List<Predicate<Long>>> predicateMap = new HashMap<>();
+ List<Predicate<Long>> predicates = predicateMap.getOrDefault(field, new ArrayList<>());
+
+ // Create a predicate to match the input query value
+ Predicate<Long> predicate = dimVal -> dimVal == inputQueryVal;
+ predicates.add(predicate);
+
+ // Put the predicates list back into the map
+ predicateMap.put(field, predicates);
+ return predicateMap;
+ }
+
+ public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory) {
+ String field;
+ Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
+ .stream()
+ .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
+
+ MetricStat metricStat = StarTreeUtils.aggregatorStatMap.get(aggregatorFactory.getClass());
+ if (metricStat != null) {
+ field = ((ValuesSourceAggregatorFactory) aggregatorFactory).getField();
+ return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
+ }
+ return false;
+ }
+
public Index index() {
return indexSettings.getIndex();
}
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index c2f8b17fcf166..5f2efe9327c96 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -77,12 +77,16 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.engine.Engine;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.DerivedFieldResolver;
import org.opensearch.index.mapper.DerivedFieldResolverFactory;
+import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.InnerHitContextBuilder;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.MatchNoneQueryBuilder;
+import org.opensearch.index.query.ParsedQuery;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryRewriteContext;
import org.opensearch.index.query.QueryShardContext;
@@ -97,11 +101,13 @@
import org.opensearch.script.ScriptService;
import org.opensearch.search.aggregations.AggregationInitializationException;
import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.InternalAggregation;
import org.opensearch.search.aggregations.InternalAggregation.ReduceContext;
import org.opensearch.search.aggregations.MultiBucketConsumerService;
import org.opensearch.search.aggregations.SearchContextAggregations;
import org.opensearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.collapse.CollapseContext;
import org.opensearch.search.deciders.ConcurrentSearchRequestDecider;
@@ -164,6 +170,7 @@
import static org.opensearch.common.unit.TimeValue.timeValueHours;
import static org.opensearch.common.unit.TimeValue.timeValueMillis;
import static org.opensearch.common.unit.TimeValue.timeValueMinutes;
+import static org.opensearch.search.internal.SearchContext.TRACK_TOTAL_HITS_DISABLED;
/**
* The main search service
@@ -1358,6 +1365,11 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.evaluateRequestShouldUseConcurrentSearch();
return;
}
+ // Can be marked false for majority cases for which star-tree cannot be used
+ // As we increment the cases where star-tree can be used, this can be set back to true
+ boolean canUseStarTree = this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
+ && context.mapperService().isCompositeIndexPresent();
+
SearchShardTarget shardTarget = context.shardTarget();
QueryShardContext queryShardContext = context.getQueryShardContext();
context.from(source.from());
@@ -1368,10 +1380,12 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.parsedQuery(queryShardContext.toQuery(source.query()));
}
if (source.postFilter() != null) {
+ canUseStarTree = false;
InnerHitContextBuilder.extractInnerHits(source.postFilter(), innerHitBuilders);
context.parsedPostFilter(queryShardContext.toQuery(source.postFilter()));
}
- if (innerHitBuilders.size() > 0) {
+ if (!innerHitBuilders.isEmpty()) {
+ canUseStarTree = false;
for (Map.Entry<String, InnerHitContextBuilder> entry : innerHitBuilders.entrySet()) {
try {
entry.getValue().build(context, context.innerHits());
@@ -1381,11 +1395,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
}
if (source.sorts() != null) {
+ canUseStarTree = false;
try {
Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(source.sorts(), context.getQueryShardContext());
- if (optionalSort.isPresent()) {
- context.sort(optionalSort.get());
- }
+ optionalSort.ifPresent(context::sort);
} catch (IOException e) {
throw new SearchException(shardTarget, "failed to create sort elements", e);
}
@@ -1399,8 +1412,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
if (source.trackTotalHitsUpTo() != null) {
context.trackTotalHitsUpTo(source.trackTotalHitsUpTo());
+ canUseStarTree = canUseStarTree && (source.trackTotalHitsUpTo() == TRACK_TOTAL_HITS_DISABLED);
}
if (source.minScore() != null) {
+ canUseStarTree = false;
context.minimumScore(source.minScore());
}
if (source.timeout() != null) {
@@ -1540,6 +1555,50 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
if (source.profile()) {
context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch()));
}
+
+ if (canUseStarTree) {
+ try {
+ if (setStarTreeQuery(context, queryShardContext, source)) {
+ logger.debug("can use star tree");
+ }
+ } catch (IOException ignored) {}
+ logger.debug("cannot use star tree");
+ }
+ }
+
+ private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source)
+ throws IOException {
+
+ if (source.aggregations() == null) {
+ return false;
+ }
+
+ // Current implementation assumes only single star-tree is supported
+ CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
+ .getCompositeFieldTypes()
+ .iterator()
+ .next();
+ CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo(
+ compositeMappedFieldType.name(),
+ compositeMappedFieldType.getCompositeIndexType()
+ );
+
+ ParsedQuery newParsedQuery = queryShardContext.toStarTreeQuery(starTree, compositeMappedFieldType, source.query(), context.query());
+ if (newParsedQuery == null) {
+ return false;
+ }
+
+ for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
+ if (!(aggregatorFactory instanceof ValuesSourceAggregatorFactory
+ && aggregatorFactory.getSubFactories().getFactories().length == 0)) {
+ return false;
+ }
+ if (queryShardContext.validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
+ return false;
+ }
+ }
+ context.parsedQuery(newParsedQuery);
+ return true;
}
/**
@@ -1699,7 +1758,7 @@ public static boolean canMatchSearchAfter(
&& minMax != null
&& primarySortField != null
&& primarySortField.missing() == null
- && Objects.equals(trackTotalHitsUpto, SearchContext.TRACK_TOTAL_HITS_DISABLED)) {
+ && Objects.equals(trackTotalHitsUpto, TRACK_TOTAL_HITS_DISABLED)) {
final Object searchAfterPrimary = searchAfter.fields[0];
if (primarySortField.order() == SortOrder.DESC) {
if (minMax.compareMin(searchAfterPrimary) > 0) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
index eeb0c606694b0..dfcb245ef3656 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
@@ -255,7 +255,7 @@ public static Builder builder() {
return new Builder();
}
- private AggregatorFactories(AggregatorFactory[] factories) {
+ public AggregatorFactories(AggregatorFactory[] factories) {
this.factories = factories;
}
@@ -661,4 +661,8 @@ public PipelineTree buildPipelineTree() {
return new PipelineTree(subTrees, aggregators);
}
}
+
+ public AggregatorFactory[] getFactories() {
+ return factories;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java
index 6cc3a78fb1e36..86fbb46a9ad3c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java
@@ -127,4 +127,8 @@ protected boolean supportsConcurrentSegmentSearch() {
public boolean evaluateChildFactories() {
return factories.allFactoriesSupportConcurrentSearch();
}
+
+ public AggregatorFactories getSubFactories() {
+ return factories;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index e58466b56df2a..adafcc82d86d4 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -32,11 +32,16 @@
package org.opensearch.search.aggregations.metrics;
import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.common.util.LongArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -46,6 +51,8 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -93,6 +100,15 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource == null) {
return LeafBucketCollector.NO_OP_COLLECTOR;
}
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ }
+ return getDefaultLeafCollector(ctx, sub);
+ }
+
+ private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
+
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
@@ -126,6 +142,56 @@ public void collect(int doc, long bucket) throws IOException {
};
}
+ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ final BigArrays bigArrays = context.bigArrays();
+ final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.SUM.getTypeName()
+ );
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(sumMetricName);
+
+ String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.VALUE_COUNT.getTypeName()
+ );
+ SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(countMetricName);
+
+ return new LeafBucketCollectorBase(sub, values) {
+ @Override
+ public void collect(int doc, long bucket) throws IOException {
+ counts = bigArrays.grow(counts, bucket + 1);
+ sums = bigArrays.grow(sums, bucket + 1);
+ compensations = bigArrays.grow(compensations, bucket + 1);
+
+ if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
+ final long valueCount = values.docValueCount();
+ counts.increment(bucket, countValues.nextValue());
+ // Compute the sum of double values with Kahan summation algorithm which is more
+ // accurate than naive summation.
+ double sum = sums.get(bucket);
+ double compensation = compensations.get(bucket);
+
+ kahanSummation.reset(sum, compensation);
+
+ for (int i = 0; i < valueCount; i++) {
+ double value = Double.longBitsToDouble(values.nextValue());
+ kahanSummation.add(value);
+ }
+
+ sums.set(bucket, kahanSummation.value());
+ compensations.set(bucket, kahanSummation.delta());
+ }
+ }
+ };
+ }
+
@Override
public double metric(long owningBucketOrd) {
if (valuesSource == null || owningBucketOrd >= sums.size()) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
index 0a09fae1eaebe..d8fe51b153819 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-class AvgAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class AvgAggregatorFactory extends ValuesSourceAggregatorFactory {
AvgAggregatorFactory(
String name,
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 8108b8a726856..bb4301c34d089 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -34,12 +34,16 @@
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PointValues;
+import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
@@ -51,6 +55,8 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Arrays;
@@ -120,6 +126,16 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
throw new CollectionTerminatedException();
}
}
+
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ }
+ return getDefaultLeafCollector(ctx, sub);
+ }
+
+ private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
+
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
final NumericDoubleValues values = MultiValueMode.MAX.select(allValues);
@@ -143,6 +159,34 @@ public void collect(int doc, long bucket) throws IOException {
};
}
+ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "max");
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+
+ final BigArrays bigArrays = context.bigArrays();
+ final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
+ return new LeafBucketCollectorBase(sub, allValues) {
+
+ @Override
+ public void collect(int doc, long bucket) throws IOException {
+ if (bucket >= maxes.size()) {
+ long from = maxes.size();
+ maxes = bigArrays.grow(maxes, bucket + 1);
+ maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
+ }
+ if (values.advanceExact(doc)) {
+ final double value = Double.longBitsToDouble(values.nextValue());
+ double max = maxes.get(bucket);
+ max = Math.max(max, value);
+ maxes.set(bucket, max);
+ }
+ }
+ };
+ }
+
@Override
public double metric(long owningBucketOrd) {
if (valuesSource == null || owningBucketOrd >= maxes.size()) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
index 4fe936c8b7797..0d537745126d3 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-class MaxAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class MaxAggregatorFactory extends ValuesSourceAggregatorFactory {
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 946057e42ac88..31e0ace591582 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -34,12 +34,16 @@
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PointValues;
+import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
@@ -51,6 +55,8 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -119,6 +125,15 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
throw new CollectionTerminatedException();
}
}
+
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ }
+ return getDefaultLeafCollector(ctx, sub);
+ }
+
+ private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
final NumericDoubleValues values = MultiValueMode.MIN.select(allValues);
@@ -138,7 +153,34 @@ public void collect(int doc, long bucket) throws IOException {
mins.set(bucket, min);
}
}
+ };
+ }
+ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "min");
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+
+ final BigArrays bigArrays = context.bigArrays();
+ final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
+ return new LeafBucketCollectorBase(sub, allValues) {
+
+ @Override
+ public void collect(int doc, long bucket) throws IOException {
+ if (bucket >= mins.size()) {
+ long from = mins.size();
+ mins = bigArrays.grow(mins, bucket + 1);
+ mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
+ }
+ if (values.advanceExact(doc)) {
+ final double value = Double.longBitsToDouble(values.nextValue());
+ double min = mins.get(bucket);
+ min = Math.min(min, value);
+ mins.set(bucket, min);
+ }
+ }
};
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
index 58fbe5edefd12..2b7c8a4cc8c9c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-class MinAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class MinAggregatorFactory extends ValuesSourceAggregatorFactory {
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index f90e5a092385f..395bf32c7ffe8 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -31,13 +31,20 @@
package org.opensearch.search.aggregations.metrics;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.opensearch.common.lucene.Lucene;
import org.opensearch.common.util.Comparators;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.sort.SortOrder;
import java.io.IOException;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
/**
* Base class to aggregate all docs into a single numeric metric value.
@@ -107,4 +114,16 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(key, lhs), metric(key, rhs), order == SortOrder.ASC);
}
}
+
+ protected StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+ SegmentReader reader = Lucene.segmentReader(ctx.reader());
+ if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) {
+ return null;
+ }
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+ StarTreeValues values = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
+ final AtomicReference<StarTreeValues> aggrVal = new AtomicReference<>(null);
+
+ return values;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 4b8e882cd69bc..b4af1d4768d11 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -32,10 +32,14 @@
package org.opensearch.search.aggregations.metrics;
import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -45,6 +49,8 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -56,13 +62,13 @@
*/
public class SumAggregator extends NumericMetricsAggregator.SingleValue {
- private final ValuesSource.Numeric valuesSource;
- private final DocValueFormat format;
+ protected final ValuesSource.Numeric valuesSource;
+ protected final DocValueFormat format;
- private DoubleArray sums;
- private DoubleArray compensations;
+ protected DoubleArray sums;
+ protected DoubleArray compensations;
- SumAggregator(
+ public SumAggregator(
String name,
ValuesSourceConfig valuesSourceConfig,
SearchContext context,
@@ -89,6 +95,15 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource == null) {
return LeafBucketCollector.NO_OP_COLLECTOR;
}
+
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ }
+ return getDefaultLeafCollector(ctx, sub);
+ }
+
+ private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
@@ -118,6 +133,40 @@ public void collect(int doc, long bucket) throws IOException {
};
}
+ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "sum");
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+
+ final BigArrays bigArrays = context.bigArrays();
+ final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+
+ return new LeafBucketCollectorBase(sub, values) {
+ @Override
+ public void collect(int doc, long bucket) throws IOException {
+ sums = bigArrays.grow(sums, bucket + 1);
+ compensations = bigArrays.grow(compensations, bucket + 1);
+
+ if (values.advanceExact(doc)) {
+ final int valuesCount = values.docValueCount();
+ double sum = sums.get(bucket);
+ double compensation = compensations.get(bucket);
+ kahanSummation.reset(sum, compensation);
+
+ for (int i = 0; i < valuesCount; i++) {
+ double value = Double.longBitsToDouble(values.nextValue());
+ kahanSummation.add(value);
+ }
+
+ compensations.set(bucket, kahanSummation.delta());
+ sums.set(bucket, kahanSummation.value());
+ }
+ }
+ };
+ }
+
@Override
public double metric(long owningBucketOrd) {
if (valuesSource == null || owningBucketOrd >= sums.size()) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
index ef9b93920ba18..e0cd44f2672a8 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
SumAggregatorFactory(
String name,
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index 6f9be06231819..7e49a4c3a1763 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -37,6 +37,10 @@
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.LongArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.MultiGeoPointValues;
import org.opensearch.index.fielddata.SortedBinaryDocValues;
import org.opensearch.search.aggregations.Aggregator;
@@ -46,6 +50,8 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -88,6 +94,12 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
final BigArrays bigArrays = context.bigArrays();
if (valuesSource instanceof ValuesSource.Numeric) {
+
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ }
+
final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx);
return new LeafBucketCollectorBase(sub, values) {
@@ -124,7 +136,29 @@ public void collect(int doc, long bucket) throws IOException {
counts.increment(bucket, values.docValueCount());
}
}
+ };
+ }
+
+ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.VALUE_COUNT.getTypeName()
+ );
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+ final BigArrays bigArrays = context.bigArrays();
+ return new LeafBucketCollectorBase(sub, values) {
+ @Override
+ public void collect(int doc, long bucket) throws IOException {
+ counts = bigArrays.grow(counts, bucket + 1);
+ if (values.advanceExact(doc)) {
+ counts.increment(bucket, values.nextValue());
+ }
+ }
};
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
index 4a04dd2e0a932..4c6fc142c182e 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
@@ -51,7 +51,7 @@
*
* @opensearch.internal
*/
-class ValueCountAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class ValueCountAggregatorFactory extends ValuesSourceAggregatorFactory {
public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(ValueCountAggregationBuilder.REGISTRY_KEY, CoreValuesSourceType.ALL_CORE, ValueCountAggregator::new, true);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSource.java b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSource.java
index 1f4dd429e094e..5732d545cb2d2 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSource.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSource.java
@@ -625,6 +625,10 @@ public SortedNumericDocValues longValues(LeafReaderContext context) {
public SortedNumericDoubleValues doubleValues(LeafReaderContext context) {
return indexFieldData.load(context).getDoubleValues();
}
+
+ public String getIndexFieldName() {
+ return indexFieldData.getFieldName();
+ }
}
/**
diff --git a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
index 69a4a5d8b6703..b19e466b081f9 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
@@ -102,4 +102,16 @@ protected abstract Aggregator doCreateInternal(
public String getStatsSubtype() {
return config.valueSourceType().typeName();
}
+
+ public String getField() {
+ return config.fieldContext().field();
+ }
+
+ public String getAggregationName() {
+ return name;
+ }
+
+ public ValuesSourceConfig getConfig() {
+ return config;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
new file mode 100644
index 0000000000000..a3aef0eab3ec9
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
@@ -0,0 +1,83 @@
+/*
+ * 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;
+
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Weight;
+import org.apache.lucene.util.Accountable;
+
+import java.io.IOException;
+
+/**
+ * Preserves star-tree queries which can be used along with original query
+ * Decides which star-tree query to use (or not) based on cost factors
+ *
+ * @opensearch.experimental
+ */
+public class OriginalOrStarTreeQuery extends Query implements Accountable {
+
+ private final StarTreeQuery starTreeQuery;
+ private final Query originalQuery;
+ private boolean starTreeQueryUsed;
+
+ public OriginalOrStarTreeQuery(StarTreeQuery starTreeQuery, Query originalQuery) {
+ this.starTreeQuery = starTreeQuery;
+ this.originalQuery = originalQuery;
+ this.starTreeQueryUsed = false;
+ }
+
+ @Override
+ public String toString(String s) {
+ return "";
+ }
+
+ @Override
+ public void visit(QueryVisitor queryVisitor) {
+
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ return true;
+ }
+
+ @Override
+ public int hashCode() {
+ return originalQuery.hashCode();
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0;
+ }
+
+ public boolean isStarTreeUsed() {
+ return starTreeQueryUsed;
+ }
+
+ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+ if (searcher.getIndexReader().hasDeletions() == false) {
+ this.starTreeQueryUsed = true;
+ return this.starTreeQuery.createWeight(searcher, scoreMode, boost);
+ } else {
+ return this.originalQuery.createWeight(searcher, scoreMode, boost);
+ }
+ }
+
+ public Query getOriginalQuery() {
+ return originalQuery;
+ }
+
+ public StarTreeQuery getStarTreeQuery() {
+ return starTreeQuery;
+ }
+}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
new file mode 100644
index 0000000000000..91ce9c06a4c60
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -0,0 +1,289 @@
+/*
+ * 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;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.DocIdSetBuilder;
+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.utils.StarTreeUtils;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.Set;
+import java.util.function.Predicate;
+
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+/**
+ * Filter operator for star tree data structure.
+ *
+ * @opensearch.experimental
+ */
+public class StarTreeFilter {
+ private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
+
+ /**
+ * Helper class to wrap the result from traversing the star tree.
+ * */
+ static class StarTreeResult {
+ final DocIdSetBuilder _matchedDocIds;
+ final Set<String> _remainingPredicateColumns;
+ final int numOfMatchedDocs;
+ final int maxMatchedDoc;
+
+ StarTreeResult(DocIdSetBuilder matchedDocIds, Set<String> remainingPredicateColumns, int numOfMatchedDocs, int maxMatchedDoc) {
+ _matchedDocIds = matchedDocIds;
+ _remainingPredicateColumns = remainingPredicateColumns;
+ this.numOfMatchedDocs = numOfMatchedDocs;
+ this.maxMatchedDoc = maxMatchedDoc;
+ }
+ }
+
+ private final StarTreeNode starTreeRoot;
+
+ Map<String, List<Predicate<Long>>> _predicateEvaluators;
+
+ DocIdSetBuilder docsWithField;
+
+ DocIdSetBuilder.BulkAdder adder;
+ Map<String, DocIdSetIterator> dimValueMap;
+
+ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, List<Predicate<Long>>> predicateEvaluators) {
+ // This filter operator does not support AND/OR/NOT operations.
+ starTreeRoot = starTreeAggrStructure.getRoot();
+ dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
+ _predicateEvaluators = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
+ // _groupByColumns = groupByColumns != null ? groupByColumns : Collections.emptyList();
+
+ // TODO : this should be the maximum number of doc values
+ docsWithField = new DocIdSetBuilder(Integer.MAX_VALUE);
+ }
+
+ /**
+ * <ul>
+ * <li>First go over the star tree and try to match as many dimensions as possible
+ * <li>For the remaining columns, use doc values indexes to match them
+ * </ul>
+ */
+ public DocIdSetIterator getStarTreeResult() throws IOException {
+ StarTreeResult starTreeResult = traverseStarTree();
+ List<DocIdSetIterator> andIterators = new ArrayList<>();
+ andIterators.add(starTreeResult._matchedDocIds.build().iterator());
+ DocIdSetIterator docIdSetIterator = andIterators.get(0);
+ // No matches, return
+ if (starTreeResult.maxMatchedDoc == -1) {
+ return docIdSetIterator;
+ }
+ int docCount = 0;
+ for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
+ // TODO : set to max value of doc values
+ logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
+ DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
+ List<Predicate<Long>> compositePredicateEvaluators = _predicateEvaluators.get(remainingPredicateColumn);
+ SortedNumericDocValues ndv = (SortedNumericDocValues) this.dimValueMap.get(remainingPredicateColumn);
+ List<Integer> docIds = new ArrayList<>();
+ while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
+ docCount++;
+ int docID = docIdSetIterator.docID();
+ if (ndv.advanceExact(docID)) {
+ final int valuesCount = ndv.docValueCount();
+ long value = ndv.nextValue();
+ for (Predicate<Long> compositePredicateEvaluator : compositePredicateEvaluators) {
+ // TODO : this might be expensive as its done against all doc values docs
+ if (compositePredicateEvaluator.test(value)) {
+ docIds.add(docID);
+ for (int i = 0; i < valuesCount - 1; i++) {
+ while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
+ docIds.add(docIdSetIterator.docID());
+ }
+ }
+ break;
+ }
+ }
+ }
+ }
+ DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size());
+ for (int docID : docIds) {
+ adder.add(docID);
+ }
+ docIdSetIterator = builder.build().iterator();
+ }
+ return docIdSetIterator;
+ }
+
+ /**
+ * Helper method to traverse the star tree, get matching documents and keep track of all the
+ * predicate dimensions that are not matched.
+ */
+ private StarTreeResult traverseStarTree() throws IOException {
+ Set<String> globalRemainingPredicateColumns = null;
+
+ StarTreeNode starTree = starTreeRoot;
+
+ List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
+
+ // Track whether we have found a leaf node added to the queue. If we have found a leaf node, and
+ // traversed to the
+ // level of the leave node, we can set globalRemainingPredicateColumns if not already set
+ // because we know the leaf
+ // node won't split further on other predicate columns.
+ boolean foundLeafNode = starTree.isLeaf();
+
+ // Use BFS to traverse the star tree
+ Queue<StarTreeNode> queue = new ArrayDeque<>();
+ queue.add(starTree);
+ int currentDimensionId = -1;
+ Set<String> remainingPredicateColumns = new HashSet<>(_predicateEvaluators.keySet());
+ if (foundLeafNode) {
+ globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
+ }
+
+ int matchedDocsCountInStarTree = 0;
+ int maxDocNum = -1;
+
+ StarTreeNode starTreeNode;
+ List<Integer> docIds = new ArrayList<>();
+ while ((starTreeNode = queue.poll()) != null) {
+ int dimensionId = starTreeNode.getDimensionId();
+ if (dimensionId > currentDimensionId) {
+ // Previous level finished
+ String dimension = dimensionNames.get(dimensionId);
+ remainingPredicateColumns.remove(dimension);
+ if (foundLeafNode && globalRemainingPredicateColumns == null) {
+ globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
+ }
+ currentDimensionId = dimensionId;
+ }
+
+ // If all predicate columns columns are matched, we can use aggregated document
+ if (remainingPredicateColumns.isEmpty()) {
+ int docId = starTreeNode.getAggregatedDocId();
+ docIds.add(docId);
+ matchedDocsCountInStarTree++;
+ maxDocNum = Math.max(docId, maxDocNum);
+ continue;
+ }
+
+ // For leaf node, because we haven't exhausted all predicate columns and group-by columns,
+ // we cannot use the aggregated document.
+ // Add the range of documents for this node to the bitmap, and keep track of the
+ // remaining predicate columns for this node
+ if (starTreeNode.isLeaf()) {
+ for (long i = starTreeNode.getStartDocId(); i < starTreeNode.getEndDocId(); i++) {
+ docIds.add((int) i);
+ matchedDocsCountInStarTree++;
+ maxDocNum = Math.max((int) i, maxDocNum);
+ }
+ continue;
+ }
+
+ // For non-leaf node, proceed to next level
+ String childDimension = dimensionNames.get(dimensionId + 1);
+
+ // Only read star-node when the dimension is not in the global remaining predicate columns
+ // because we cannot use star-node in such cases
+ StarTreeNode starNode = null;
+ if ((globalRemainingPredicateColumns == null || !globalRemainingPredicateColumns.contains(childDimension))) {
+ starNode = starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true);
+ }
+
+ if (remainingPredicateColumns.contains(childDimension)) {
+ // Have predicates on the next level, add matching nodes to the queue
+
+ // Calculate the matching dictionary ids for the child dimension
+ int numChildren = starTreeNode.getNumChildren();
+
+ // If number of matching dictionary ids is large, use scan instead of binary search
+
+ Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
+
+ // When the star-node exists, and the number of matching doc ids is more than or equal to
+ // the number of non-star child nodes, check if all the child nodes match the predicate,
+ // and use the star-node if so
+ if (starNode != null) {
+ List<StarTreeNode> matchingChildNodes = new ArrayList<>();
+ boolean findLeafChildNode = false;
+ while (childrenIterator.hasNext()) {
+ StarTreeNode childNode = childrenIterator.next();
+ List<Predicate<Long>> predicates = _predicateEvaluators.get(childDimension);
+ for (Predicate<Long> predicate : predicates) {
+ long val = childNode.getDimensionValue();
+ if (predicate.test(val)) {
+ matchingChildNodes.add(childNode);
+ findLeafChildNode |= childNode.isLeaf();
+ break;
+ }
+ }
+ }
+ if (matchingChildNodes.size() == numChildren - 1) {
+ // All the child nodes (except for the star-node) match the predicate, use the star-node
+ queue.add(starNode);
+ foundLeafNode |= starNode.isLeaf();
+ } else {
+ // Some child nodes do not match the predicate, use the matching child nodes
+ queue.addAll(matchingChildNodes);
+ foundLeafNode |= findLeafChildNode;
+ }
+ } else {
+ // Cannot use the star-node, use the matching child nodes
+ while (childrenIterator.hasNext()) {
+ StarTreeNode childNode = childrenIterator.next();
+ List<Predicate<Long>> predicates = _predicateEvaluators.get(childDimension);
+ for (Predicate<Long> predicate : predicates) {
+ if (predicate.test(childNode.getDimensionValue())) {
+ queue.add(childNode);
+ foundLeafNode |= childNode.isLeaf();
+ break;
+ }
+ }
+ }
+ }
+ } else {
+ // No predicate on the next level
+ if (starNode != null) {
+ // Star-node exists, use it
+ queue.add(starNode);
+ foundLeafNode |= starNode.isLeaf();
+ } else {
+ // Star-node does not exist or cannot be used, add all non-star nodes to the queue
+ Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
+ while (childrenIterator.hasNext()) {
+ StarTreeNode childNode = childrenIterator.next();
+ if (childNode.getDimensionValue() != StarTreeUtils.ALL) {
+ queue.add(childNode);
+ foundLeafNode |= childNode.isLeaf();
+ }
+ }
+ }
+ }
+ }
+
+ adder = docsWithField.grow(docIds.size());
+ for (int id : docIds) {
+ adder.add(id);
+ }
+ return new StarTreeResult(
+ docsWithField,
+ globalRemainingPredicateColumns != null ? globalRemainingPredicateColumns : Collections.emptySet(),
+ matchedDocsCountInStarTree,
+ maxDocNum
+ );
+ }
+}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
new file mode 100644
index 0000000000000..0c24e2f469909
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -0,0 +1,116 @@
+/*
+ * 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;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.apache.lucene.search.ConstantScoreScorer;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.Weight;
+import org.apache.lucene.util.Accountable;
+import org.opensearch.common.lucene.Lucene;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Predicate;
+
+/**
+ * Query class for querying star tree data structure.
+ *
+ * @opensearch.experimental
+ */
+public class StarTreeQuery extends Query implements Accountable {
+
+ /**
+ * Star tree field info
+ * This is used to get the star tree data structure
+ */
+ CompositeIndexFieldInfo starTree;
+
+ /**
+ * Map of field name to a list of predicates to be applied on that field
+ * This is used to filter the data based on the predicates
+ */
+ Map<String, List<Predicate<Long>>> compositePredicateMap;
+
+ public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, List<Predicate<Long>>> compositePredicateMap) {
+ this.starTree = starTree;
+ this.compositePredicateMap = compositePredicateMap;
+ }
+
+ @Override
+ public String toString(String field) {
+ return null;
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {
+ visitor.visitLeaf(this);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ return sameClassAs(obj);
+ }
+
+ @Override
+ public int hashCode() {
+ return classHash();
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return 0;
+ }
+
+ @Override
+ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+ return new ConstantScoreWeight(this, boost) {
+ @Override
+ public Scorer scorer(LeafReaderContext context) throws IOException {
+ SegmentReader reader = Lucene.segmentReader(context.reader());
+
+ // We get the 'CompositeIndexReader' instance so that we can get StarTreeValues
+ if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) return null;
+
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+ List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+ StarTreeValues starTreeValues = null;
+ if (compositeIndexFields != null && !compositeIndexFields.isEmpty()) {
+ starTreeValues = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
+ } else {
+ return null;
+ }
+
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, compositePredicateMap);
+ DocIdSetIterator result = filter.getStarTreeResult();
+ return new ConstantScoreScorer(this, score(), scoreMode, result);
+ }
+
+ @Override
+ public boolean isCacheable(LeafReaderContext ctx) {
+ return false;
+ }
+ };
+ }
+
+ public CompositeIndexFieldInfo getStarTree() {
+ return starTree;
+ }
+}
diff --git a/server/src/main/java/org/opensearch/search/startree/package-info.java b/server/src/main/java/org/opensearch/search/startree/package-info.java
new file mode 100644
index 0000000000000..601a588e54e69
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/startree/package-info.java
@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+/** Star Tree query classes */
+package org.opensearch.search.startree;
diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
index d35fc6b111c9f..2aaa815662427 100644
--- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
+++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
@@ -307,7 +307,7 @@ public void testStarTreeDocValuesWithDeletions() throws IOException {
directory.close();
}
- private XContentBuilder getExpandedMapping() throws IOException {
+ public static XContentBuilder getExpandedMapping() throws IOException {
return topMapping(b -> {
b.startObject("composite");
b.startObject("startree");
@@ -361,13 +361,14 @@ private XContentBuilder getExpandedMapping() throws IOException {
});
}
- private XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
+ private static XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
XContentBuilder builder = XContentFactory.jsonBuilder().startObject().startObject("_doc");
buildFields.accept(builder);
return builder.endObject().endObject();
}
- private void createMapperService(XContentBuilder builder) throws IOException {
+
+ public static MapperService createMapperService(XContentBuilder builder) throws IOException {
Settings settings = Settings.builder()
.put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
@@ -377,7 +378,7 @@ private void createMapperService(XContentBuilder builder) throws IOException {
.build();
IndexMetadata indexMetadata = IndexMetadata.builder("test").settings(settings).putMapping(builder.toString()).build();
IndicesModule indicesModule = new IndicesModule(Collections.emptyList());
- mapperService = MapperTestUtils.newMapperServiceWithHelperAnalyzer(
+ MapperService mapperService = MapperTestUtils.newMapperServiceWithHelperAnalyzer(
new NamedXContentRegistry(ClusterModule.getNamedXWriteables()),
createTempDir(),
settings,
@@ -385,5 +386,6 @@ private void createMapperService(XContentBuilder builder) throws IOException {
"test"
);
mapperService.merge(indexMetadata, MapperService.MergeReason.INDEX_TEMPLATE);
+ return mapperService;
}
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
new file mode 100644
index 0000000000000..515382967f4a1
--- /dev/null
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -0,0 +1,212 @@
+/*
+ * 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.aggregations.startree;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.lucene99.Lucene99Codec;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.opensearch.common.lucene.Lucene;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
+import org.opensearch.index.codec.composite.composite99.Composite99Codec;
+import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.MapperService;
+import org.opensearch.index.mapper.NumberFieldMapper;
+import org.opensearch.search.aggregations.AggregationBuilder;
+import org.opensearch.search.aggregations.AggregatorTestCase;
+import org.opensearch.search.aggregations.InternalAggregation;
+import org.opensearch.search.aggregations.metrics.AvgAggregationBuilder;
+import org.opensearch.search.aggregations.metrics.InternalAvg;
+import org.opensearch.search.aggregations.metrics.InternalMax;
+import org.opensearch.search.aggregations.metrics.InternalMin;
+import org.opensearch.search.aggregations.metrics.InternalSum;
+import org.opensearch.search.aggregations.metrics.InternalValueCount;
+import org.opensearch.search.aggregations.metrics.MaxAggregationBuilder;
+import org.opensearch.search.aggregations.metrics.MinAggregationBuilder;
+import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
+import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+import java.util.function.Predicate;
+
+import static org.opensearch.search.aggregations.AggregationBuilders.avg;
+import static org.opensearch.search.aggregations.AggregationBuilders.count;
+import static org.opensearch.search.aggregations.AggregationBuilders.max;
+import static org.opensearch.search.aggregations.AggregationBuilders.min;
+import static org.opensearch.search.aggregations.AggregationBuilders.sum;
+import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS;
+
+public class MetricAggregatorTests extends AggregatorTestCase {
+
+ private static final String FIELD_NAME = "field";
+ private static final NumberFieldMapper.NumberType DEFAULT_FIELD_TYPE = NumberFieldMapper.NumberType.LONG;
+ private static final MappedFieldType DEFAULT_MAPPED_FIELD = new NumberFieldMapper.NumberFieldType(FIELD_NAME, DEFAULT_FIELD_TYPE);
+
+ @Before
+ public void setup() {
+ FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+ }
+
+ @After
+ public void teardown() throws IOException {
+ FeatureFlags.initializeFeatureFlags(Settings.EMPTY);
+ }
+
+ protected Codec getCodec() {
+ final Logger testLogger = LogManager.getLogger(MetricAggregatorTests.class);
+ MapperService mapperService;
+ try {
+ mapperService = StarTreeDocValuesFormatTests.createMapperService(StarTreeDocValuesFormatTests.getExpandedMapping());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ return new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
+ }
+
+ public void testStarTreeDocValues() throws IOException {
+ Directory directory = newDirectory();
+ IndexWriterConfig conf = newIndexWriterConfig(null);
+ conf.setCodec(getCodec());
+ conf.setMergePolicy(newLogMergePolicy());
+ RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
+ Document doc = new Document();
+ doc.add(new SortedNumericDocValuesField("sndv", 1));
+ doc.add(new SortedNumericDocValuesField("dv", 1));
+ doc.add(new SortedNumericDocValuesField("field", 1));
+ iw.addDocument(doc);
+ doc = new Document();
+ doc.add(new SortedNumericDocValuesField("sndv", 1));
+ doc.add(new SortedNumericDocValuesField("dv", 2));
+ doc.add(new SortedNumericDocValuesField("field", 3));
+ iw.addDocument(doc);
+ doc = new Document();
+ iw.forceMerge(1);
+ doc.add(new SortedNumericDocValuesField("sndv", 2));
+ doc.add(new SortedNumericDocValuesField("dv", 4));
+ doc.add(new SortedNumericDocValuesField("field", 6));
+ iw.addDocument(doc);
+ doc = new Document();
+ doc.add(new SortedNumericDocValuesField("sndv", 3));
+ doc.add(new SortedNumericDocValuesField("dv", 6));
+ doc.add(new SortedNumericDocValuesField("field", 9));
+ iw.addDocument(doc);
+ iw.forceMerge(1);
+ iw.close();
+
+ DirectoryReader ir = DirectoryReader.open(directory);
+ initValuesSourceRegistry();
+ assertEquals(ir.leaves().size(), 1);
+ LeafReaderContext context = ir.leaves().get(0);
+
+ SegmentReader reader = Lucene.segmentReader(context.reader());
+ IndexSearcher indexSearcher = newSearcher(reader, true, true);
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+ List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+
+ CompositeIndexFieldInfo starTree = compositeIndexFields.get(0);
+
+ SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME);
+ MaxAggregationBuilder maxAggregationBuilder = max("_name").field(FIELD_NAME);
+ MinAggregationBuilder minAggregationBuilder = min("_name").field(FIELD_NAME);
+ ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
+ AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
+
+ // match-all query
+ Query defaultQeury = new MatchAllDocsQuery();
+ StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates
+ testCase(indexSearcher, defaultQeury, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue), 19);
+ testCase(indexSearcher, defaultQeury, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue), 9);
+ testCase(indexSearcher, defaultQeury, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue), 1);
+ testCase(
+ indexSearcher,
+ defaultQeury,
+ starTreeQuery,
+ valueCountAggregationBuilder,
+ verifyAggregation(InternalValueCount::getValue),
+ 4
+ );
+ testCase(indexSearcher, defaultQeury, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue), 4.75);
+
+ // numeric-terms query
+ defaultQeury = new TermQuery(new Term("sndv", "1"));
+ Map<String, List<Predicate<Long>>> compositePredicateMap = new HashMap<>();
+ compositePredicateMap.put("sndv", List.of(dimVal -> dimVal == 1));
+ starTreeQuery = new StarTreeQuery(starTree, compositePredicateMap);
+ testCase(indexSearcher, defaultQeury, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue), 4);
+ testCase(indexSearcher, defaultQeury, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue), 3);
+ testCase(indexSearcher, defaultQeury, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue), 1);
+ testCase(
+ indexSearcher,
+ defaultQeury,
+ starTreeQuery,
+ valueCountAggregationBuilder,
+ verifyAggregation(InternalValueCount::getValue),
+ 2
+ );
+ testCase(indexSearcher, defaultQeury, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue), 2);
+
+ ir.close();
+ directory.close();
+ }
+
+ private <T extends AggregationBuilder, V extends InternalAggregation, N extends Number> void testCase(
+ IndexSearcher searcher,
+ Query defaultQuery,
+ StarTreeQuery starTreeQuery,
+ T builder,
+ BiConsumer<V, N> verify,
+ N expectedValue
+ ) throws IOException {
+ OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery);
+ V aggregation = searchAndReduceStarTree(
+ createIndexSettings(),
+ searcher,
+ originalOrStarTreeQuery,
+ builder,
+ DEFAULT_MAX_BUCKETS,
+ false,
+ DEFAULT_MAPPED_FIELD
+ );
+ verify.accept(aggregation, expectedValue);
+ }
+
+ <T, R extends Number> BiConsumer<T, R> verifyAggregation(Function<T, R> valueExtractor) {
+ return (aggregation, expectedValue) -> assertEquals(
+ expectedValue.doubleValue(),
+ valueExtractor.apply(aggregation).doubleValue(),
+ 0f
+ );
+ }
+}
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 4abd7fbea9cff..9d816e864c50b 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -37,6 +37,7 @@
import org.apache.lucene.document.Field;
import org.apache.lucene.document.InetAddressPoint;
import org.apache.lucene.document.LatLonDocValuesField;
+import org.apache.lucene.document.LongField;
import org.apache.lucene.document.SortedNumericDocValuesField;
import org.apache.lucene.document.SortedSetDocValuesField;
import org.apache.lucene.document.StoredField;
@@ -141,6 +142,7 @@
import org.opensearch.search.internal.ContextIndexSearcher;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.lookup.SearchLookup;
+import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.test.InternalAggregationTestCase;
import org.opensearch.test.OpenSearchTestCase;
import org.junit.After;
@@ -650,6 +652,74 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer);
return internalAgg;
}
+ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduceStarTree(
+ IndexSettings indexSettings,
+ IndexSearcher searcher,
+ Query query,
+ AggregationBuilder builder,
+ int maxBucket,
+ boolean hasNested,
+ MappedFieldType... fieldTypes
+ ) throws IOException {
+ final IndexReaderContext ctx = searcher.getTopReaderContext();
+ final PipelineTree pipelines = builder.buildPipelineTree();
+ List<InternalAggregation> aggs = new ArrayList<>();
+ if (hasNested) {
+ query = Queries.filtered(query, Queries.newNonNestedFilter());
+ }
+
+ MultiBucketConsumer bucketConsumer = new MultiBucketConsumer(
+ maxBucket,
+ new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
+ );
+ C root = createAggregator(query, builder, searcher, bucketConsumer, fieldTypes);
+
+ if (randomBoolean() && searcher.getIndexReader().leaves().size() > 0) {
+ assertTrue(ctx instanceof LeafReaderContext);
+ final LeafReaderContext compCTX = (LeafReaderContext) ctx;
+ final int size = compCTX.leaves().size();
+ final ShardSearcher[] subSearchers = new ShardSearcher[size];
+ for (int searcherIDX = 0; searcherIDX < subSearchers.length; searcherIDX++) {
+ final LeafReaderContext leave = compCTX.leaves().get(searcherIDX);
+ subSearchers[searcherIDX] = new ShardSearcher(leave, compCTX);
+ }
+ for (ShardSearcher subSearcher : subSearchers) {
+ MultiBucketConsumer shardBucketConsumer = new MultiBucketConsumer(
+ maxBucket,
+ new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
+ );
+ C a = createAggregator(query, builder, subSearcher, indexSettings, shardBucketConsumer, fieldTypes);
+ a.preCollection();
+ Weight weight = subSearcher.createWeight(query, ScoreMode.COMPLETE, 1f);
+
+ assertTrue(weight.getQuery() instanceof StarTreeQuery);
+ subSearcher.search(weight, a);
+ a.postCollection();
+ aggs.add(a.buildTopLevel());
+ }
+ } else {
+ root.preCollection();
+ searcher.search(query, root);
+ root.postCollection();
+ aggs.add(root.buildTopLevel());
+ }
+
+ MultiBucketConsumer reduceBucketConsumer = new MultiBucketConsumer(
+ maxBucket,
+ new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
+ );
+ InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forFinalReduction(
+ root.context().bigArrays(),
+ getMockScriptService(),
+ reduceBucketConsumer,
+ pipelines
+ );
+
+ @SuppressWarnings("unchecked")
+ A internalAgg = (A) aggs.get(0).reduce(aggs, context);
+ doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer);
+ return internalAgg;
+ }
protected void doAssertReducedMultiBucketConsumer(Aggregation agg, MultiBucketConsumerService.MultiBucketConsumer bucketConsumer) {
InternalAggregationTestCase.assertMultiBucketConsumer(agg, bucketConsumer);
From a624206b7979086716a0af46ea14db8d792d40fa Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 26 Aug 2024 22:58:26 -0700
Subject: [PATCH 02/35] using bst as filter
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/startree/node/StarTreeNode.java | 2 +
.../index/query/QueryShardContext.java | 28 +--
.../aggregations/metrics/AvgAggregator.java | 3 +-
.../aggregations/metrics/MaxAggregator.java | 3 +-
.../aggregations/metrics/MinAggregator.java | 3 +-
.../aggregations/metrics/SumAggregator.java | 3 +-
.../startree/OriginalOrStarTreeQuery.java | 21 +-
.../search/startree/StarTreeFilter.java | 215 ++++++------------
.../search/startree/StarTreeQuery.java | 37 ++-
.../startree/MetricAggregatorTests.java | 149 ++++++------
.../aggregations/AggregatorTestCase.java | 2 +-
11 files changed, 204 insertions(+), 262 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
index fce3e30e9ebf6..1f73da8cd66b0 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
@@ -124,4 +124,6 @@ public interface StarTreeNode {
* @throws IOException if an I/O error occurs while retrieving the children iterator
*/
Iterator<? extends StarTreeNode> getChildrenIterator() throws IOException;
+
+ Iterator<? extends StarTreeNode> getChildrenIteratorForRange(int min, int max) throws IOException;
}
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index 894f467a8f6fc..db48626dbea41 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -85,6 +85,7 @@
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.lookup.SearchLookup;
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.transport.RemoteClusterAware;
@@ -600,21 +601,21 @@ public ParsedQuery toStarTreeQuery(
QueryBuilder queryBuilder,
Query query
) {
- Map<String, List<Predicate<Long>>> predicateMap;
+ Map<String, List<StarTreeFilter.Range>> queryMap;
if (queryBuilder == null) {
- predicateMap = null;
+ queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
.stream()
.map(Dimension::getField)
.collect(Collectors.toList());
- predicateMap = getStarTreePredicates(queryBuilder, supportedDimensions);
+ queryMap = getStarTreePredicates(queryBuilder, supportedDimensions);
} else {
return null;
}
- StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, predicateMap);
+ StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, queryMap);
OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, query);
return new ParsedQuery(originalOrStarTreeQuery);
}
@@ -624,24 +625,23 @@ public ParsedQuery toStarTreeQuery(
* @param queryBuilder
* @return predicates to match
*/
- private Map<String, List<Predicate<Long>>> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
+ private Map<String, List<StarTreeFilter.Range>> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
String field = tq.fieldName();
- if (supportedDimensions.contains(field) == false) {
+
+ if (!supportedDimensions.contains(field)) {
throw new IllegalArgumentException("unsupported field in star-tree");
}
+
long inputQueryVal = Long.parseLong(tq.value().toString());
- // Get or create the list of predicates for the given field
- Map<String, List<Predicate<Long>>> predicateMap = new HashMap<>();
- List<Predicate<Long>> predicates = predicateMap.getOrDefault(field, new ArrayList<>());
+ // Create a Range with the same min and max value for the exact match for TermQuery
+ StarTreeFilter.Range range = new StarTreeFilter.Range(inputQueryVal, inputQueryVal);
- // Create a predicate to match the input query value
- Predicate<Long> predicate = dimVal -> dimVal == inputQueryVal;
- predicates.add(predicate);
+ // Create a map with the field and the range
+ Map<String, List<StarTreeFilter.Range>> predicateMap = new HashMap<>();
+ predicateMap.computeIfAbsent(field, k -> new ArrayList<>()).add(range);
- // Put the predicates list back into the map
- predicateMap.put(field, predicates);
return predicateMap;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index adafcc82d86d4..fcae269c11cca 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -34,6 +34,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
@@ -181,7 +182,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valueCount; i++) {
- double value = Double.longBitsToDouble(values.nextValue());
+ double value = NumericUtils.sortableLongToDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index bb4301c34d089..af1bf61428c24 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -38,6 +38,7 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
@@ -178,7 +179,7 @@ public void collect(int doc, long bucket) throws IOException {
maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = Double.longBitsToDouble(values.nextValue());
+ final double value = NumericUtils.sortableLongToDouble(values.nextValue());
double max = maxes.get(bucket);
max = Math.max(max, value);
maxes.set(bucket, max);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 31e0ace591582..864e0e9e4c34b 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -38,6 +38,7 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
@@ -175,7 +176,7 @@ public void collect(int doc, long bucket) throws IOException {
mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = Double.longBitsToDouble(values.nextValue());
+ final double value = NumericUtils.sortableLongToDouble(values.nextValue());
double min = mins.get(bucket);
min = Math.min(min, value);
mins.set(bucket, min);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index b4af1d4768d11..fb7d3bc7679f1 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -34,6 +34,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
@@ -156,7 +157,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valuesCount; i++) {
- double value = Double.longBitsToDouble(values.nextValue());
+ double value = NumericUtils.sortableLongToDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
index a3aef0eab3ec9..34044d1c44bc7 100644
--- a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
@@ -13,9 +13,9 @@
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Weight;
-import org.apache.lucene.util.Accountable;
import java.io.IOException;
+import java.util.Objects;
/**
* Preserves star-tree queries which can be used along with original query
@@ -23,7 +23,7 @@
*
* @opensearch.experimental
*/
-public class OriginalOrStarTreeQuery extends Query implements Accountable {
+public class OriginalOrStarTreeQuery extends Query {
private final StarTreeQuery starTreeQuery;
private final Query originalQuery;
@@ -37,27 +37,24 @@ public OriginalOrStarTreeQuery(StarTreeQuery starTreeQuery, Query originalQuery)
@Override
public String toString(String s) {
- return "";
+ return originalQuery.toString(s);
}
@Override
- public void visit(QueryVisitor queryVisitor) {
-
- }
+ public void visit(QueryVisitor queryVisitor) {}
@Override
public boolean equals(Object o) {
- return true;
+ return sameClassAs(o) && equalsTo(getClass().cast(o));
}
- @Override
- public int hashCode() {
- return originalQuery.hashCode();
+ private boolean equalsTo(OriginalOrStarTreeQuery other) {
+ return starTreeQuery.equals(other.starTreeQuery) && originalQuery.equals(other.originalQuery);
}
@Override
- public long ramBytesUsed() {
- return 0;
+ public int hashCode() {
+ return Objects.hash(classHash(), starTreeQuery, originalQuery, starTreeQuery);
}
public boolean isStarTreeUsed() {
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 91ce9c06a4c60..df54e2c9314c2 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -27,7 +27,6 @@
import java.util.Map;
import java.util.Queue;
import java.util.Set;
-import java.util.function.Predicate;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
@@ -39,38 +38,15 @@
public class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
- /**
- * Helper class to wrap the result from traversing the star tree.
- * */
- static class StarTreeResult {
- final DocIdSetBuilder _matchedDocIds;
- final Set<String> _remainingPredicateColumns;
- final int numOfMatchedDocs;
- final int maxMatchedDoc;
-
- StarTreeResult(DocIdSetBuilder matchedDocIds, Set<String> remainingPredicateColumns, int numOfMatchedDocs, int maxMatchedDoc) {
- _matchedDocIds = matchedDocIds;
- _remainingPredicateColumns = remainingPredicateColumns;
- this.numOfMatchedDocs = numOfMatchedDocs;
- this.maxMatchedDoc = maxMatchedDoc;
- }
- }
-
private final StarTreeNode starTreeRoot;
-
- Map<String, List<Predicate<Long>>> _predicateEvaluators;
-
+ Map<String, List<Range>> _predicateEvaluators;
DocIdSetBuilder docsWithField;
-
- DocIdSetBuilder.BulkAdder adder;
Map<String, DocIdSetIterator> dimValueMap;
- public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, List<Predicate<Long>>> predicateEvaluators) {
- // This filter operator does not support AND/OR/NOT operations.
+ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, List<Range>> predicateEvaluators) {
starTreeRoot = starTreeAggrStructure.getRoot();
dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
_predicateEvaluators = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
- // _groupByColumns = groupByColumns != null ? groupByColumns : Collections.emptyList();
// TODO : this should be the maximum number of doc values
docsWithField = new DocIdSetBuilder(Integer.MAX_VALUE);
@@ -87,38 +63,41 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
List<DocIdSetIterator> andIterators = new ArrayList<>();
andIterators.add(starTreeResult._matchedDocIds.build().iterator());
DocIdSetIterator docIdSetIterator = andIterators.get(0);
+
// No matches, return
if (starTreeResult.maxMatchedDoc == -1) {
return docIdSetIterator;
}
- int docCount = 0;
+
for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
- // TODO : set to max value of doc values
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
- List<Predicate<Long>> compositePredicateEvaluators = _predicateEvaluators.get(remainingPredicateColumn);
SortedNumericDocValues ndv = (SortedNumericDocValues) this.dimValueMap.get(remainingPredicateColumn);
List<Integer> docIds = new ArrayList<>();
+ List<Range> queryRanges = _predicateEvaluators.get(remainingPredicateColumn); // Get the list of min-max ranges for this field
+
while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
- docCount++;
int docID = docIdSetIterator.docID();
if (ndv.advanceExact(docID)) {
final int valuesCount = ndv.docValueCount();
- long value = ndv.nextValue();
- for (Predicate<Long> compositePredicateEvaluator : compositePredicateEvaluators) {
- // TODO : this might be expensive as its done against all doc values docs
- if (compositePredicateEvaluator.test(value)) {
- docIds.add(docID);
- for (int i = 0; i < valuesCount - 1; i++) {
- while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
- docIds.add(docIdSetIterator.docID());
- }
+ for (int i = 0; i < valuesCount; i++) {
+ double value = ndv.nextValue();
+ // Check if the value falls within any of the ranges in the query
+ boolean matchFound = false;
+ for (Range range : queryRanges) {
+ if (value >= range.getMin() && value <= range.getMax()) {
+ docIds.add(docID);
+ matchFound = true;
+ break;
}
+ }
+ if (matchFound) {
break;
}
}
}
}
+
DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size());
for (int docID : docIds) {
adder.add(docID);
@@ -133,46 +112,27 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
* predicate dimensions that are not matched.
*/
private StarTreeResult traverseStarTree() throws IOException {
- Set<String> globalRemainingPredicateColumns = null;
-
- StarTreeNode starTree = starTreeRoot;
-
- List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
-
- // Track whether we have found a leaf node added to the queue. If we have found a leaf node, and
- // traversed to the
- // level of the leave node, we can set globalRemainingPredicateColumns if not already set
- // because we know the leaf
- // node won't split further on other predicate columns.
- boolean foundLeafNode = starTree.isLeaf();
-
- // Use BFS to traverse the star tree
Queue<StarTreeNode> queue = new ArrayDeque<>();
- queue.add(starTree);
- int currentDimensionId = -1;
- Set<String> remainingPredicateColumns = new HashSet<>(_predicateEvaluators.keySet());
- if (foundLeafNode) {
- globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
- }
+ queue.add(starTreeRoot);
+ Set<String> globalRemainingPredicateColumns = null;
int matchedDocsCountInStarTree = 0;
int maxDocNum = -1;
-
- StarTreeNode starTreeNode;
List<Integer> docIds = new ArrayList<>();
- while ((starTreeNode = queue.poll()) != null) {
+ boolean foundLeafNode = false;
+ int currentDimensionId = -1;
+ Set<String> remainingPredicateColumns = new HashSet<>(_predicateEvaluators.keySet());
+ List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
+
+ while (!queue.isEmpty()) {
+ StarTreeNode starTreeNode = queue.poll();
int dimensionId = starTreeNode.getDimensionId();
if (dimensionId > currentDimensionId) {
- // Previous level finished
String dimension = dimensionNames.get(dimensionId);
remainingPredicateColumns.remove(dimension);
- if (foundLeafNode && globalRemainingPredicateColumns == null) {
- globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
- }
currentDimensionId = dimensionId;
}
- // If all predicate columns columns are matched, we can use aggregated document
if (remainingPredicateColumns.isEmpty()) {
int docId = starTreeNode.getAggregatedDocId();
docIds.add(docId);
@@ -181,10 +141,6 @@ private StarTreeResult traverseStarTree() throws IOException {
continue;
}
- // For leaf node, because we haven't exhausted all predicate columns and group-by columns,
- // we cannot use the aggregated document.
- // Add the range of documents for this node to the bitmap, and keep track of the
- // remaining predicate columns for this node
if (starTreeNode.isLeaf()) {
for (long i = starTreeNode.getStartDocId(); i < starTreeNode.getEndDocId(); i++) {
docIds.add((int) i);
@@ -194,91 +150,29 @@ private StarTreeResult traverseStarTree() throws IOException {
continue;
}
- // For non-leaf node, proceed to next level
String childDimension = dimensionNames.get(dimensionId + 1);
-
- // Only read star-node when the dimension is not in the global remaining predicate columns
- // because we cannot use star-node in such cases
- StarTreeNode starNode = null;
- if ((globalRemainingPredicateColumns == null || !globalRemainingPredicateColumns.contains(childDimension))) {
- starNode = starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true);
- }
-
if (remainingPredicateColumns.contains(childDimension)) {
- // Have predicates on the next level, add matching nodes to the queue
+ List<Range> queryRanges = _predicateEvaluators.get(childDimension);
- // Calculate the matching dictionary ids for the child dimension
- int numChildren = starTreeNode.getNumChildren();
+ for (Range range : queryRanges) {
+ Iterator rangeIterator = starTreeNode.getChildrenIteratorForRange((int) range.getMin(), (int) range.getMax());
- // If number of matching dictionary ids is large, use scan instead of binary search
-
- Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
-
- // When the star-node exists, and the number of matching doc ids is more than or equal to
- // the number of non-star child nodes, check if all the child nodes match the predicate,
- // and use the star-node if so
- if (starNode != null) {
- List<StarTreeNode> matchingChildNodes = new ArrayList<>();
- boolean findLeafChildNode = false;
- while (childrenIterator.hasNext()) {
- StarTreeNode childNode = childrenIterator.next();
- List<Predicate<Long>> predicates = _predicateEvaluators.get(childDimension);
- for (Predicate<Long> predicate : predicates) {
- long val = childNode.getDimensionValue();
- if (predicate.test(val)) {
- matchingChildNodes.add(childNode);
- findLeafChildNode |= childNode.isLeaf();
- break;
- }
- }
- }
- if (matchingChildNodes.size() == numChildren - 1) {
- // All the child nodes (except for the star-node) match the predicate, use the star-node
- queue.add(starNode);
- foundLeafNode |= starNode.isLeaf();
- } else {
- // Some child nodes do not match the predicate, use the matching child nodes
- queue.addAll(matchingChildNodes);
- foundLeafNode |= findLeafChildNode;
- }
- } else {
- // Cannot use the star-node, use the matching child nodes
- while (childrenIterator.hasNext()) {
- StarTreeNode childNode = childrenIterator.next();
- List<Predicate<Long>> predicates = _predicateEvaluators.get(childDimension);
- for (Predicate<Long> predicate : predicates) {
- if (predicate.test(childNode.getDimensionValue())) {
- queue.add(childNode);
- foundLeafNode |= childNode.isLeaf();
- break;
- }
- }
+ while (rangeIterator.hasNext()) {
+ StarTreeNode matchingChildNode = (StarTreeNode) rangeIterator.next();
+ queue.add(matchingChildNode);
+ foundLeafNode |= matchingChildNode.isLeaf();
}
}
} else {
- // No predicate on the next level
- if (starNode != null) {
- // Star-node exists, use it
- queue.add(starNode);
- foundLeafNode |= starNode.isLeaf();
- } else {
- // Star-node does not exist or cannot be used, add all non-star nodes to the queue
- Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
- while (childrenIterator.hasNext()) {
- StarTreeNode childNode = childrenIterator.next();
- if (childNode.getDimensionValue() != StarTreeUtils.ALL) {
- queue.add(childNode);
- foundLeafNode |= childNode.isLeaf();
- }
- }
- }
+ queue.add(starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true));
}
}
- adder = docsWithField.grow(docIds.size());
+ DocIdSetBuilder.BulkAdder adder = docsWithField.grow(docIds.size());
for (int id : docIds) {
adder.add(id);
}
+
return new StarTreeResult(
docsWithField,
globalRemainingPredicateColumns != null ? globalRemainingPredicateColumns : Collections.emptySet(),
@@ -286,4 +180,39 @@ private StarTreeResult traverseStarTree() throws IOException {
maxDocNum
);
}
+
+ /**
+ * Helper class to wrap the result from traversing the star tree.
+ * */
+ private static class StarTreeResult {
+ final DocIdSetBuilder _matchedDocIds;
+ final Set<String> _remainingPredicateColumns;
+ final int numOfMatchedDocs;
+ final int maxMatchedDoc;
+
+ StarTreeResult(DocIdSetBuilder matchedDocIds, Set<String> remainingPredicateColumns, int numOfMatchedDocs, int maxMatchedDoc) {
+ _matchedDocIds = matchedDocIds;
+ _remainingPredicateColumns = remainingPredicateColumns;
+ this.numOfMatchedDocs = numOfMatchedDocs;
+ this.maxMatchedDoc = maxMatchedDoc;
+ }
+ }
+
+ public static class Range {
+ long min;
+ long max;
+
+ public Range(long min, long max) {
+ this.min = min;
+ this.max = max;
+ }
+
+ public long getMin() {
+ return min;
+ }
+
+ public long getMax() {
+ return max;
+ }
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index 0c24e2f469909..7661af6548e76 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -19,7 +19,6 @@
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Weight;
-import org.apache.lucene.util.Accountable;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite.CompositeIndexReader;
@@ -28,14 +27,14 @@
import java.io.IOException;
import java.util.List;
import java.util.Map;
-import java.util.function.Predicate;
+import java.util.Objects;
/**
* Query class for querying star tree data structure.
*
* @opensearch.experimental
*/
-public class StarTreeQuery extends Query implements Accountable {
+public class StarTreeQuery extends Query {
/**
* Star tree field info
@@ -44,39 +43,37 @@ public class StarTreeQuery extends Query implements Accountable {
CompositeIndexFieldInfo starTree;
/**
- * Map of field name to a list of predicates to be applied on that field
- * This is used to filter the data based on the predicates
+ * Map of field name to a value to be queried for that field
+ * This is used to filter the data based on the query
*/
- Map<String, List<Predicate<Long>>> compositePredicateMap;
+ Map<String, List<StarTreeFilter.Range>> queryMap;
- public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, List<Predicate<Long>>> compositePredicateMap) {
+ public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, List<StarTreeFilter.Range>> queryMap) {
this.starTree = starTree;
- this.compositePredicateMap = compositePredicateMap;
+ this.queryMap = queryMap;
}
@Override
- public String toString(String field) {
- return null;
- }
+ public void visit(QueryVisitor visitor) {}
@Override
- public void visit(QueryVisitor visitor) {
- visitor.visitLeaf(this);
+ public boolean equals(Object obj) {
+ return sameClassAs(obj) && equalsTo(getClass().cast(obj));
}
- @Override
- public boolean equals(Object obj) {
- return sameClassAs(obj);
+ private boolean equalsTo(StarTreeQuery other) {
+ return starTree.equals(other.starTree) && queryMap != null && queryMap.equals(other.queryMap);
}
@Override
public int hashCode() {
- return classHash();
+ return Objects.hash(classHash(), starTree, queryMap);
}
@Override
- public long ramBytesUsed() {
- return 0;
+ public String toString(String field) {
+ // Does not implements a user-readable toString
+ return null;
}
@Override
@@ -98,7 +95,7 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
return null;
}
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, compositePredicateMap);
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, queryMap);
DocIdSetIterator result = filter.getStarTreeResult();
return new ConstantScoreScorer(this, score(), scoreMode, result);
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 515382967f4a1..6cefcfe512a31 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -18,11 +18,9 @@
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
-import org.apache.lucene.index.Term;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
-import org.apache.lucene.search.TermQuery;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.opensearch.common.lucene.Lucene;
@@ -49,17 +47,17 @@
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQuery;
import org.junit.After;
import org.junit.Before;
import java.io.IOException;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Random;
import java.util.function.BiConsumer;
import java.util.function.Function;
-import java.util.function.Predicate;
import static org.opensearch.search.aggregations.AggregationBuilders.avg;
import static org.opensearch.search.aggregations.AggregationBuilders.count;
@@ -101,28 +99,30 @@ public void testStarTreeDocValues() throws IOException {
conf.setCodec(getCodec());
conf.setMergePolicy(newLogMergePolicy());
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
- Document doc = new Document();
- doc.add(new SortedNumericDocValuesField("sndv", 1));
- doc.add(new SortedNumericDocValuesField("dv", 1));
- doc.add(new SortedNumericDocValuesField("field", 1));
- iw.addDocument(doc);
- doc = new Document();
- doc.add(new SortedNumericDocValuesField("sndv", 1));
- doc.add(new SortedNumericDocValuesField("dv", 2));
- doc.add(new SortedNumericDocValuesField("field", 3));
- iw.addDocument(doc);
- doc = new Document();
- iw.forceMerge(1);
- doc.add(new SortedNumericDocValuesField("sndv", 2));
- doc.add(new SortedNumericDocValuesField("dv", 4));
- doc.add(new SortedNumericDocValuesField("field", 6));
- iw.addDocument(doc);
- doc = new Document();
- doc.add(new SortedNumericDocValuesField("sndv", 3));
- doc.add(new SortedNumericDocValuesField("dv", 6));
- doc.add(new SortedNumericDocValuesField("field", 9));
- iw.addDocument(doc);
- iw.forceMerge(1);
+
+ Random random = new Random();
+ int totalDocs = 100;
+ final String SNDV = "sndv";
+ final String DV = "dv";
+
+ // Index 100 random documents
+ for (int i = 0; i < totalDocs; i++) {
+ Document doc = new Document();
+ if (random.nextBoolean()) {
+ doc.add(new SortedNumericDocValuesField(SNDV, random.nextInt(10) - 5)); // Random long between -5 and 4
+ }
+ if (random.nextBoolean()) {
+ doc.add(new SortedNumericDocValuesField(DV, random.nextInt(20) - 10)); // Random long between -10 and 9
+ }
+ if (random.nextBoolean()) {
+ doc.add(new SortedNumericDocValuesField(FIELD_NAME, random.nextInt(50))); // Random long between 0 and 49
+ }
+ iw.addDocument(doc);
+ }
+
+ if (randomBoolean()) {
+ iw.forceMerge(1);
+ }
iw.close();
DirectoryReader ir = DirectoryReader.open(directory);
@@ -131,7 +131,7 @@ public void testStarTreeDocValues() throws IOException {
LeafReaderContext context = ir.leaves().get(0);
SegmentReader reader = Lucene.segmentReader(context.reader());
- IndexSearcher indexSearcher = newSearcher(reader, true, true);
+ IndexSearcher indexSearcher = newSearcher(reader, false, false);
CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
@@ -144,53 +144,57 @@ public void testStarTreeDocValues() throws IOException {
AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
// match-all query
- Query defaultQeury = new MatchAllDocsQuery();
+ Query defaultQuery = new MatchAllDocsQuery();
StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates
- testCase(indexSearcher, defaultQeury, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue), 19);
- testCase(indexSearcher, defaultQeury, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue), 9);
- testCase(indexSearcher, defaultQeury, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue), 1);
- testCase(
- indexSearcher,
- defaultQeury,
- starTreeQuery,
- valueCountAggregationBuilder,
- verifyAggregation(InternalValueCount::getValue),
- 4
- );
- testCase(indexSearcher, defaultQeury, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue), 4.75);
-
+ testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
// numeric-terms query
- defaultQeury = new TermQuery(new Term("sndv", "1"));
- Map<String, List<Predicate<Long>>> compositePredicateMap = new HashMap<>();
- compositePredicateMap.put("sndv", List.of(dimVal -> dimVal == 1));
- starTreeQuery = new StarTreeQuery(starTree, compositePredicateMap);
- testCase(indexSearcher, defaultQeury, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue), 4);
- testCase(indexSearcher, defaultQeury, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue), 3);
- testCase(indexSearcher, defaultQeury, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue), 1);
- testCase(
- indexSearcher,
- defaultQeury,
- starTreeQuery,
- valueCountAggregationBuilder,
- verifyAggregation(InternalValueCount::getValue),
- 2
- );
- testCase(indexSearcher, defaultQeury, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue), 2);
-
+ for (int cases = 0; cases < 100; cases++) {
+ Map<String, List<StarTreeFilter.Range>> queryMap;
+ String queryField;
+ long queryValue;
+ if (randomBoolean()) {
+ queryField = SNDV;
+ queryValue = random.nextInt(10) - 5;
+ } else {
+ queryField = DV;
+ queryValue = random.nextInt(20) - 10;
+ }
+ if (queryValue == -1) {
+ continue; // -1 is encoded as null
+ }
+ defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
+ queryMap = Map.of(queryField, List.of(new StarTreeFilter.Range(queryValue, queryValue)));
+ starTreeQuery = new StarTreeQuery(starTree, queryMap);
+
+ testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
+ testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
+ testCase(
+ indexSearcher,
+ defaultQuery,
+ starTreeQuery,
+ valueCountAggregationBuilder,
+ verifyAggregation(InternalValueCount::getValue)
+ );
+ testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
+ }
ir.close();
directory.close();
}
- private <T extends AggregationBuilder, V extends InternalAggregation, N extends Number> void testCase(
+ private <T extends AggregationBuilder, V extends InternalAggregation, R extends Number> void testCase(
IndexSearcher searcher,
Query defaultQuery,
StarTreeQuery starTreeQuery,
T builder,
- BiConsumer<V, N> verify,
- N expectedValue
+ BiConsumer<V, V> verify
) throws IOException {
OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery);
- V aggregation = searchAndReduceStarTree(
+ V starTreeAggregation = searchAndReduceStarTree(
createIndexSettings(),
searcher,
originalOrStarTreeQuery,
@@ -199,14 +203,23 @@ private <T extends AggregationBuilder, V extends InternalAggregation, N extends
false,
DEFAULT_MAPPED_FIELD
);
- verify.accept(aggregation, expectedValue);
+ V expectedAggregation = searchAndReduceStarTree(
+ createIndexSettings(),
+ searcher,
+ defaultQuery,
+ builder,
+ DEFAULT_MAX_BUCKETS,
+ false,
+ DEFAULT_MAPPED_FIELD
+ );
+ verify.accept(expectedAggregation, starTreeAggregation);
}
- <T, R extends Number> BiConsumer<T, R> verifyAggregation(Function<T, R> valueExtractor) {
- return (aggregation, expectedValue) -> assertEquals(
- expectedValue.doubleValue(),
- valueExtractor.apply(aggregation).doubleValue(),
- 0f
+ <T, R extends Number> BiConsumer<T, T> verifyAggregation(Function<T, R> valueExtractor) {
+ return (expectedAggregation, actualAggregation) -> assertEquals(
+ valueExtractor.apply(expectedAggregation).doubleValue(),
+ valueExtractor.apply(actualAggregation).doubleValue(),
+ 0.0f
);
}
}
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 9d816e864c50b..6f9b6330ac95e 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -692,7 +692,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
a.preCollection();
Weight weight = subSearcher.createWeight(query, ScoreMode.COMPLETE, 1f);
- assertTrue(weight.getQuery() instanceof StarTreeQuery);
+// assertTrue(weight.getQuery() instanceof StarTreeQuery);
subSearcher.search(weight, a);
a.postCollection();
aggs.add(a.buildTopLevel());
From da9e40bec4537acb308b8d40d330f41941fbc47a Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Thu, 29 Aug 2024 01:05:21 -0600
Subject: [PATCH 03/35] Refactoring & bug fixes
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../java/org/opensearch/common/util/FeatureFlags.java | 2 +-
.../search/aggregations/metrics/MaxAggregator.java | 7 ++++++-
.../search/aggregations/metrics/MinAggregator.java | 7 ++++++-
.../search/aggregations/metrics/SumAggregator.java | 7 ++++++-
.../support/ValuesSourceAggregatorFactory.java | 8 --------
.../aggregations/startree/MetricAggregatorTests.java | 1 -
6 files changed, 19 insertions(+), 13 deletions(-)
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index 6df68013a8119..e663d8429da13 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index af1bf61428c24..356100d9caf11 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -43,6 +43,7 @@
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -164,7 +165,11 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
throws IOException {
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "max");
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.MAX.getTypeName()
+ );
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
final BigArrays bigArrays = context.bigArrays();
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 864e0e9e4c34b..6ecc57a6b23ad 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -43,6 +43,7 @@
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -161,7 +162,11 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
throws IOException {
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "min");
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.MIN.getTypeName()
+ );
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
final BigArrays bigArrays = context.bigArrays();
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index fb7d3bc7679f1..4e8bc0a415c0f 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -39,6 +39,7 @@
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -138,7 +139,11 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
throws IOException {
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, "sum");
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.SUM.getTypeName()
+ );
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
final BigArrays bigArrays = context.bigArrays();
diff --git a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
index b19e466b081f9..d862b2c2784de 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java
@@ -106,12 +106,4 @@ public String getStatsSubtype() {
public String getField() {
return config.fieldContext().field();
}
-
- public String getAggregationName() {
- return name;
- }
-
- public ValuesSourceConfig getConfig() {
- return config;
- }
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 6cefcfe512a31..6ac6e7d41e255 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -127,7 +127,6 @@ public void testStarTreeDocValues() throws IOException {
DirectoryReader ir = DirectoryReader.open(directory);
initValuesSourceRegistry();
- assertEquals(ir.leaves().size(), 1);
LeafReaderContext context = ir.leaves().get(0);
SegmentReader reader = Lucene.segmentReader(context.reader());
From 810b2f317bf5acabd3095df1c73708d5298b9d2c Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Sat, 31 Aug 2024 11:24:48 -0600
Subject: [PATCH 04/35] change filter logic and aggragtor factory
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeUtils.java | 24 ----
.../index/query/QueryShardContext.java | 28 ++---
.../org/opensearch/search/SearchService.java | 7 +-
.../aggregations/AggregatorFactories.java | 2 +-
.../metrics/AvgAggregatorFactory.java | 9 +-
.../metrics/MaxAggregatorFactory.java | 9 +-
.../metrics/MetricAggregatorFactory.java | 37 ++++++
.../metrics/MinAggregatorFactory.java | 9 +-
.../metrics/SumAggregatorFactory.java | 9 +-
.../metrics/ValueCountAggregatorFactory.java | 9 +-
.../search/startree/StarTreeFilter.java | 108 ++++++++----------
.../search/startree/StarTreeQuery.java | 8 +-
.../startree/MetricAggregatorTests.java | 20 +++-
13 files changed, 153 insertions(+), 126 deletions(-)
create mode 100644 server/src/main/java/org/opensearch/search/aggregations/metrics/MetricAggregatorFactory.java
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
index 201bb895ee49a..2aae0d4ca7e29 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeUtils.java
@@ -12,18 +12,10 @@
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.VectorEncoding;
import org.apache.lucene.index.VectorSimilarityFunction;
-import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo;
-import org.opensearch.search.aggregations.metrics.AvgAggregatorFactory;
-import org.opensearch.search.aggregations.metrics.MaxAggregatorFactory;
-import org.opensearch.search.aggregations.metrics.MinAggregatorFactory;
-import org.opensearch.search.aggregations.metrics.SumAggregatorFactory;
-import org.opensearch.search.aggregations.metrics.ValueCountAggregatorFactory;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import java.util.Collections;
import java.util.List;
-import java.util.Map;
/**
* Util class for building star tree
@@ -46,22 +38,6 @@ private StarTreeUtils() {}
*/
public static final String METRIC_SUFFIX = "metric";
- /**
- * Map to associate star-tree supported AggregatorFactory classes with their corresponding MetricStat
- */
- public static final Map<Class<? extends ValuesSourceAggregatorFactory>, MetricStat> aggregatorStatMap = Map.of(
- SumAggregatorFactory.class,
- MetricStat.SUM,
- MaxAggregatorFactory.class,
- MetricStat.MAX,
- MinAggregatorFactory.class,
- MetricStat.MIN,
- ValueCountAggregatorFactory.class,
- MetricStat.VALUE_COUNT,
- AvgAggregatorFactory.class,
- MetricStat.AVG
- );
-
/**
* Returns the full field name for a dimension in the star-tree index.
*
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index db48626dbea41..2539775c3b541 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -60,7 +60,6 @@
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.Metric;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.IndexFieldData;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.ContentPath;
@@ -80,17 +79,15 @@
import org.opensearch.script.ScriptFactory;
import org.opensearch.script.ScriptService;
import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
import org.opensearch.search.aggregations.support.AggregationUsageService;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.lookup.SearchLookup;
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.transport.RemoteClusterAware;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
@@ -601,8 +598,7 @@ public ParsedQuery toStarTreeQuery(
QueryBuilder queryBuilder,
Query query
) {
- Map<String, List<StarTreeFilter.Range>> queryMap;
-
+ Map<String, Long> queryMap;
if (queryBuilder == null) {
queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
@@ -625,23 +621,17 @@ public ParsedQuery toStarTreeQuery(
* @param queryBuilder
* @return predicates to match
*/
- private Map<String, List<StarTreeFilter.Range>> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
+ private Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
String field = tq.fieldName();
-
if (!supportedDimensions.contains(field)) {
throw new IllegalArgumentException("unsupported field in star-tree");
}
-
long inputQueryVal = Long.parseLong(tq.value().toString());
- // Create a Range with the same min and max value for the exact match for TermQuery
- StarTreeFilter.Range range = new StarTreeFilter.Range(inputQueryVal, inputQueryVal);
-
- // Create a map with the field and the range
- Map<String, List<StarTreeFilter.Range>> predicateMap = new HashMap<>();
- predicateMap.computeIfAbsent(field, k -> new ArrayList<>()).add(range);
-
+ // Create a map with the field and the value
+ Map<String, Long> predicateMap = new HashMap<>();
+ predicateMap.put(field, inputQueryVal);
return predicateMap;
}
@@ -651,9 +641,9 @@ public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType composite
.stream()
.collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
- MetricStat metricStat = StarTreeUtils.aggregatorStatMap.get(aggregatorFactory.getClass());
- if (metricStat != null) {
- field = ((ValuesSourceAggregatorFactory) aggregatorFactory).getField();
+ if (aggregatorFactory instanceof MetricAggregatorFactory) {
+ MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
+ field = ((MetricAggregatorFactory) aggregatorFactory).getField();
return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
}
return false;
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 5f2efe9327c96..950dad4cd720e 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -1412,7 +1412,9 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
if (source.trackTotalHitsUpTo() != null) {
context.trackTotalHitsUpTo(source.trackTotalHitsUpTo());
- canUseStarTree = canUseStarTree && (source.trackTotalHitsUpTo() == TRACK_TOTAL_HITS_DISABLED);
+ if (source.trackTotalHitsUpTo() != TRACK_TOTAL_HITS_DISABLED) {
+ canUseStarTree = false;
+ }
}
if (source.minScore() != null) {
canUseStarTree = false;
@@ -1560,9 +1562,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
try {
if (setStarTreeQuery(context, queryShardContext, source)) {
logger.debug("can use star tree");
+ } else {
+ logger.debug("cannot use star tree");
}
} catch (IOException ignored) {}
- logger.debug("cannot use star tree");
}
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
index dfcb245ef3656..ab3d78032dc54 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
@@ -255,7 +255,7 @@ public static Builder builder() {
return new Builder();
}
- public AggregatorFactories(AggregatorFactory[] factories) {
+ AggregatorFactories(AggregatorFactory[] factories) {
this.factories = factories;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
index d8fe51b153819..57389f19b4577 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregatorFactory.java
@@ -32,13 +32,13 @@
package org.opensearch.search.aggregations.metrics;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.aggregations.AggregatorFactories;
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.CardinalityUpperBound;
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.internal.SearchContext;
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-public class AvgAggregatorFactory extends ValuesSourceAggregatorFactory {
+class AvgAggregatorFactory extends MetricAggregatorFactory {
AvgAggregatorFactory(
String name,
@@ -65,6 +65,11 @@ public class AvgAggregatorFactory extends ValuesSourceAggregatorFactory {
super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
}
+ @Override
+ public MetricStat getMetricStat() {
+ return MetricStat.AVG;
+ }
+
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
AvgAggregationBuilder.REGISTRY_KEY,
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
index 0d537745126d3..c0ee471c87f29 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregatorFactory.java
@@ -32,13 +32,13 @@
package org.opensearch.search.aggregations.metrics;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.aggregations.AggregatorFactories;
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.CardinalityUpperBound;
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.internal.SearchContext;
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-public class MaxAggregatorFactory extends ValuesSourceAggregatorFactory {
+class MaxAggregatorFactory extends MetricAggregatorFactory {
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
@@ -74,6 +74,11 @@ static void registerAggregators(ValuesSourceRegistry.Builder builder) {
super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
}
+ @Override
+ public MetricStat getMetricStat() {
+ return MetricStat.MAX;
+ }
+
@Override
protected Aggregator createUnmapped(SearchContext searchContext, Aggregator parent, Map<String, Object> metadata) throws IOException {
return new MaxAggregator(name, config, searchContext, parent, metadata);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MetricAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MetricAggregatorFactory.java
new file mode 100644
index 0000000000000..0ac630cf051d3
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MetricAggregatorFactory.java
@@ -0,0 +1,37 @@
+/*
+ * 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.aggregations.metrics;
+
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceConfig;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Extending ValuesSourceAggregatorFactory for aggregation factories supported by star-tree implementation
+ */
+public abstract class MetricAggregatorFactory extends ValuesSourceAggregatorFactory {
+ public MetricAggregatorFactory(
+ String name,
+ ValuesSourceConfig config,
+ QueryShardContext queryShardContext,
+ AggregatorFactory parent,
+ AggregatorFactories.Builder subFactoriesBuilder,
+ Map<String, Object> metadata
+ ) throws IOException {
+ super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
+ }
+
+ public abstract MetricStat getMetricStat();
+}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
index 2b7c8a4cc8c9c..44c0d9d7d11eb 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregatorFactory.java
@@ -32,13 +32,13 @@
package org.opensearch.search.aggregations.metrics;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.aggregations.AggregatorFactories;
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.CardinalityUpperBound;
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.internal.SearchContext;
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-public class MinAggregatorFactory extends ValuesSourceAggregatorFactory {
+class MinAggregatorFactory extends MetricAggregatorFactory {
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
@@ -74,6 +74,11 @@ static void registerAggregators(ValuesSourceRegistry.Builder builder) {
super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
}
+ @Override
+ public MetricStat getMetricStat() {
+ return MetricStat.MIN;
+ }
+
@Override
protected Aggregator createUnmapped(SearchContext searchContext, Aggregator parent, Map<String, Object> metadata) throws IOException {
return new MinAggregator(name, config, searchContext, parent, metadata);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
index e0cd44f2672a8..e2e25a8c25a87 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
@@ -32,13 +32,13 @@
package org.opensearch.search.aggregations.metrics;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.aggregations.AggregatorFactories;
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.CardinalityUpperBound;
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.internal.SearchContext;
@@ -52,7 +52,7 @@
*
* @opensearch.internal
*/
-public class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
+class SumAggregatorFactory extends MetricAggregatorFactory {
SumAggregatorFactory(
String name,
@@ -65,6 +65,11 @@ public class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
}
+ @Override
+ public MetricStat getMetricStat() {
+ return MetricStat.SUM;
+ }
+
static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(
SumAggregationBuilder.REGISTRY_KEY,
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
index 4c6fc142c182e..0c82279484461 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
@@ -32,13 +32,13 @@
package org.opensearch.search.aggregations.metrics;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.aggregations.AggregatorFactories;
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.CardinalityUpperBound;
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.internal.SearchContext;
@@ -51,7 +51,7 @@
*
* @opensearch.internal
*/
-public class ValueCountAggregatorFactory extends ValuesSourceAggregatorFactory {
+class ValueCountAggregatorFactory extends MetricAggregatorFactory {
public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(ValueCountAggregationBuilder.REGISTRY_KEY, CoreValuesSourceType.ALL_CORE, ValueCountAggregator::new, true);
@@ -68,6 +68,11 @@ public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata);
}
+ @Override
+ public MetricStat getMetricStat() {
+ return MetricStat.VALUE_COUNT;
+ }
+
@Override
protected Aggregator createUnmapped(SearchContext searchContext, Aggregator parent, Map<String, Object> metadata) throws IOException {
return new ValueCountAggregator(name, config, searchContext, parent, metadata);
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index df54e2c9314c2..4ff9885cf2563 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -39,14 +39,16 @@ public class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
private final StarTreeNode starTreeRoot;
- Map<String, List<Range>> _predicateEvaluators;
+ Map<String, Long> queryMap;
DocIdSetBuilder docsWithField;
+ DocIdSetBuilder.BulkAdder adder;
Map<String, DocIdSetIterator> dimValueMap;
- public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, List<Range>> predicateEvaluators) {
+ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> predicateEvaluators) {
+ // This filter operator does not support AND/OR/NOT operations as of now.
starTreeRoot = starTreeAggrStructure.getRoot();
dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
- _predicateEvaluators = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
+ queryMap = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
// TODO : this should be the maximum number of doc values
docsWithField = new DocIdSetBuilder(Integer.MAX_VALUE);
@@ -68,36 +70,27 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
if (starTreeResult.maxMatchedDoc == -1) {
return docIdSetIterator;
}
-
for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
SortedNumericDocValues ndv = (SortedNumericDocValues) this.dimValueMap.get(remainingPredicateColumn);
List<Integer> docIds = new ArrayList<>();
- List<Range> queryRanges = _predicateEvaluators.get(remainingPredicateColumn); // Get the list of min-max ranges for this field
+ long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly
while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
int docID = docIdSetIterator.docID();
if (ndv.advanceExact(docID)) {
final int valuesCount = ndv.docValueCount();
for (int i = 0; i < valuesCount; i++) {
- double value = ndv.nextValue();
- // Check if the value falls within any of the ranges in the query
- boolean matchFound = false;
- for (Range range : queryRanges) {
- if (value >= range.getMin() && value <= range.getMax()) {
- docIds.add(docID);
- matchFound = true;
- break;
- }
- }
- if (matchFound) {
+ long value = ndv.nextValue();
+ // Directly compare value with queryValue
+ if (value == queryValue) {
+ docIds.add(docID);
break;
}
}
}
}
-
DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size());
for (int docID : docIds) {
adder.add(docID);
@@ -112,24 +105,30 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
* predicate dimensions that are not matched.
*/
private StarTreeResult traverseStarTree() throws IOException {
- Queue<StarTreeNode> queue = new ArrayDeque<>();
- queue.add(starTreeRoot);
Set<String> globalRemainingPredicateColumns = null;
-
+ StarTreeNode starTree = starTreeRoot;
+ List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
+ boolean foundLeafNode = starTree.isLeaf();
+ Queue<StarTreeNode> queue = new ArrayDeque<>();
+ queue.add(starTree);
+ int currentDimensionId = -1;
+ Set<String> remainingPredicateColumns = new HashSet<>(queryMap.keySet());
+ if (foundLeafNode) {
+ globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
+ }
int matchedDocsCountInStarTree = 0;
int maxDocNum = -1;
+ StarTreeNode starTreeNode;
List<Integer> docIds = new ArrayList<>();
- boolean foundLeafNode = false;
- int currentDimensionId = -1;
- Set<String> remainingPredicateColumns = new HashSet<>(_predicateEvaluators.keySet());
- List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
- while (!queue.isEmpty()) {
- StarTreeNode starTreeNode = queue.poll();
+ while ((starTreeNode = queue.poll()) != null) {
int dimensionId = starTreeNode.getDimensionId();
if (dimensionId > currentDimensionId) {
String dimension = dimensionNames.get(dimensionId);
remainingPredicateColumns.remove(dimension);
+ if (foundLeafNode && globalRemainingPredicateColumns == null) {
+ globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
+ }
currentDimensionId = dimensionId;
}
@@ -151,28 +150,39 @@ private StarTreeResult traverseStarTree() throws IOException {
}
String childDimension = dimensionNames.get(dimensionId + 1);
- if (remainingPredicateColumns.contains(childDimension)) {
- List<Range> queryRanges = _predicateEvaluators.get(childDimension);
-
- for (Range range : queryRanges) {
- Iterator rangeIterator = starTreeNode.getChildrenIteratorForRange((int) range.getMin(), (int) range.getMax());
+ StarTreeNode starNode = null;
+ if (globalRemainingPredicateColumns == null || !globalRemainingPredicateColumns.contains(childDimension)) {
+ starNode = starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true);
+ }
- while (rangeIterator.hasNext()) {
- StarTreeNode matchingChildNode = (StarTreeNode) rangeIterator.next();
- queue.add(matchingChildNode);
- foundLeafNode |= matchingChildNode.isLeaf();
- }
+ if (remainingPredicateColumns.contains(childDimension)) {
+ long queryValue = queryMap.get(childDimension); // Get the query value directly from the map
+ StarTreeNode matchingChild = starTreeNode.getChildForDimensionValue(queryValue, false);
+ if (matchingChild != null) {
+ queue.add(matchingChild);
+ foundLeafNode |= matchingChild.isLeaf();
}
} else {
- queue.add(starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true));
+ if (starNode != null) {
+ queue.add(starNode);
+ foundLeafNode |= starNode.isLeaf();
+ } else {
+ Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
+ while (childrenIterator.hasNext()) {
+ StarTreeNode childNode = childrenIterator.next();
+ if (childNode.getDimensionValue() != StarTreeUtils.ALL) {
+ queue.add(childNode);
+ foundLeafNode |= childNode.isLeaf();
+ }
+ }
+ }
}
}
- DocIdSetBuilder.BulkAdder adder = docsWithField.grow(docIds.size());
+ adder = docsWithField.grow(docIds.size());
for (int id : docIds) {
adder.add(id);
}
-
return new StarTreeResult(
docsWithField,
globalRemainingPredicateColumns != null ? globalRemainingPredicateColumns : Collections.emptySet(),
@@ -184,7 +194,7 @@ private StarTreeResult traverseStarTree() throws IOException {
/**
* Helper class to wrap the result from traversing the star tree.
* */
- private static class StarTreeResult {
+ static class StarTreeResult {
final DocIdSetBuilder _matchedDocIds;
final Set<String> _remainingPredicateColumns;
final int numOfMatchedDocs;
@@ -197,22 +207,4 @@ private static class StarTreeResult {
this.maxMatchedDoc = maxMatchedDoc;
}
}
-
- public static class Range {
- long min;
- long max;
-
- public Range(long min, long max) {
- this.min = min;
- this.max = max;
- }
-
- public long getMin() {
- return min;
- }
-
- public long getMax() {
- return max;
- }
- }
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index 7661af6548e76..1f4cdaddcfebf 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -46,9 +46,9 @@ public class StarTreeQuery extends Query {
* Map of field name to a value to be queried for that field
* This is used to filter the data based on the query
*/
- Map<String, List<StarTreeFilter.Range>> queryMap;
+ Map<String, Long> queryMap;
- public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, List<StarTreeFilter.Range>> queryMap) {
+ public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
this.starTree = starTree;
this.queryMap = queryMap;
}
@@ -82,10 +82,6 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
@Override
public Scorer scorer(LeafReaderContext context) throws IOException {
SegmentReader reader = Lucene.segmentReader(context.reader());
-
- // We get the 'CompositeIndexReader' instance so that we can get StarTreeValues
- if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) return null;
-
CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
StarTreeValues starTreeValues = null;
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 6ac6e7d41e255..9f6c2aa764bff 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -47,12 +47,12 @@
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQuery;
import org.junit.After;
import org.junit.Before;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
@@ -100,26 +100,34 @@ public void testStarTreeDocValues() throws IOException {
conf.setMergePolicy(newLogMergePolicy());
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
+ List<Document> documents = new ArrayList<>();
Random random = new Random();
int totalDocs = 100;
final String SNDV = "sndv";
final String DV = "dv";
+ int val;
// Index 100 random documents
for (int i = 0; i < totalDocs; i++) {
Document doc = new Document();
if (random.nextBoolean()) {
- doc.add(new SortedNumericDocValuesField(SNDV, random.nextInt(10) - 5)); // Random long between -5 and 4
+ val = random.nextInt(10) - 5;
+ val = val == -1 ? 1 : val; // -1 is encoded as null
+ doc.add(new SortedNumericDocValuesField(SNDV, val)); // Random long between -5 and 4
}
if (random.nextBoolean()) {
- doc.add(new SortedNumericDocValuesField(DV, random.nextInt(20) - 10)); // Random long between -10 and 9
+ val = random.nextInt(20) - 10;
+ val = val == -1 ? 1 : val; // -1 is encoded as null
+ doc.add(new SortedNumericDocValuesField(DV, val)); // Random long between -10 and 9
}
if (random.nextBoolean()) {
doc.add(new SortedNumericDocValuesField(FIELD_NAME, random.nextInt(50))); // Random long between 0 and 49
}
iw.addDocument(doc);
+ documents.add(doc);
}
+ iw.forceMerge(1);
if (randomBoolean()) {
iw.forceMerge(1);
}
@@ -151,8 +159,8 @@ public void testStarTreeDocValues() throws IOException {
testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
// numeric-terms query
- for (int cases = 0; cases < 100; cases++) {
- Map<String, List<StarTreeFilter.Range>> queryMap;
+ for (int cases = 0; cases < 1; cases++) {
+ Map<String, Long> queryMap;
String queryField;
long queryValue;
if (randomBoolean()) {
@@ -166,7 +174,7 @@ public void testStarTreeDocValues() throws IOException {
continue; // -1 is encoded as null
}
defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
- queryMap = Map.of(queryField, List.of(new StarTreeFilter.Range(queryValue, queryValue)));
+ queryMap = Map.of(queryField, queryValue);
starTreeQuery = new StarTreeQuery(starTree, queryMap);
testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
From 274f0bfa19650e9df10b6e8a6a3ecf909470488e Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 00:29:36 -0700
Subject: [PATCH 05/35] Rebasing changes with indexing / file-formats
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/startree/node/StarTreeNode.java | 2 --
.../aggregations/metrics/AvgAggregator.java | 4 +--
.../aggregations/metrics/MaxAggregator.java | 2 +-
.../aggregations/metrics/MinAggregator.java | 2 +-
.../aggregations/metrics/SumAggregator.java | 2 +-
.../metrics/ValueCountAggregator.java | 2 +-
.../search/startree/StarTreeFilter.java | 29 +++++++++++++------
.../search/startree/StarTreeQuery.java | 3 +-
.../startree/MetricAggregatorTests.java | 14 ++-------
9 files changed, 31 insertions(+), 29 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
index 1f73da8cd66b0..fce3e30e9ebf6 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/node/StarTreeNode.java
@@ -124,6 +124,4 @@ public interface StarTreeNode {
* @throws IOException if an I/O error occurs while retrieving the children iterator
*/
Iterator<? extends StarTreeNode> getChildrenIterator() throws IOException;
-
- Iterator<? extends StarTreeNode> getChildrenIteratorForRange(int min, int max) throws IOException;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index fcae269c11cca..a4338f6eb5812 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -155,14 +155,14 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.SUM.getTypeName()
);
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(sumMetricName);
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
starTree.getField(),
fieldName,
MetricStat.VALUE_COUNT.getTypeName()
);
- SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(countMetricName);
+ SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
return new LeafBucketCollectorBase(sub, values) {
@Override
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 356100d9caf11..bd511a9f2205d 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -170,7 +170,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.MAX.getTypeName()
);
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 6ecc57a6b23ad..e2a1bb053f937 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -167,7 +167,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.MIN.getTypeName()
);
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 4e8bc0a415c0f..82bdf3030b3f3 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -144,7 +144,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.SUM.getTypeName()
);
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index 7e49a4c3a1763..b0cc37d1703fd 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -148,7 +148,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.VALUE_COUNT.getTypeName()
);
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocValuesIteratorMap().get(metricName);
+ SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
return new LeafBucketCollectorBase(sub, values) {
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 4ff9885cf2563..b2838ad788da6 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -13,6 +13,7 @@
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.DocIdSetBuilder;
+import org.opensearch.index.compositeindex.datacube.Dimension;
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.utils.StarTreeUtils;
@@ -27,6 +28,7 @@
import java.util.Map;
import java.util.Queue;
import java.util.Set;
+import java.util.stream.Collectors;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
@@ -38,16 +40,19 @@
public class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
- private final StarTreeNode starTreeRoot;
+ // private final StarTreeNode starTreeRoot;
Map<String, Long> queryMap;
DocIdSetBuilder docsWithField;
DocIdSetBuilder.BulkAdder adder;
- Map<String, DocIdSetIterator> dimValueMap;
+ // Map<String, DocIdSetIterator> dimValueMap;
+ StarTreeValues starTreeValues;
+ List<Dimension> dimensions;
public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> predicateEvaluators) {
// This filter operator does not support AND/OR/NOT operations as of now.
- starTreeRoot = starTreeAggrStructure.getRoot();
- dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
+ starTreeValues = starTreeAggrStructure;
+ // starTreeRoot = starTreeAggrStructure.getRoot();
+ // dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
queryMap = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
// TODO : this should be the maximum number of doc values
@@ -73,7 +78,9 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
- SortedNumericDocValues ndv = (SortedNumericDocValues) this.dimValueMap.get(remainingPredicateColumn);
+ SortedNumericDocValues ndv = (SortedNumericDocValues) this.starTreeValues.getDimensionDocIdSetIterator(
+ remainingPredicateColumn
+ );
List<Integer> docIds = new ArrayList<>();
long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly
@@ -106,8 +113,12 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
*/
private StarTreeResult traverseStarTree() throws IOException {
Set<String> globalRemainingPredicateColumns = null;
- StarTreeNode starTree = starTreeRoot;
- List<String> dimensionNames = new ArrayList<>(dimValueMap.keySet());
+ StarTreeNode starTree = starTreeValues.getRoot();
+ List<String> dimensionNames = starTreeValues.getStarTreeField()
+ .getDimensionsOrder()
+ .stream()
+ .map(Dimension::getField)
+ .collect(Collectors.toList());
boolean foundLeafNode = starTree.isLeaf();
Queue<StarTreeNode> queue = new ArrayDeque<>();
queue.add(starTree);
@@ -152,12 +163,12 @@ private StarTreeResult traverseStarTree() throws IOException {
String childDimension = dimensionNames.get(dimensionId + 1);
StarTreeNode starNode = null;
if (globalRemainingPredicateColumns == null || !globalRemainingPredicateColumns.contains(childDimension)) {
- starNode = starTreeNode.getChildForDimensionValue(StarTreeUtils.ALL, true);
+ starNode = starTreeNode.getChildStarNode();
}
if (remainingPredicateColumns.contains(childDimension)) {
long queryValue = queryMap.get(childDimension); // Get the query value directly from the map
- StarTreeNode matchingChild = starTreeNode.getChildForDimensionValue(queryValue, false);
+ StarTreeNode matchingChild = starTreeNode.getChildForDimensionValue(queryValue);
if (matchingChild != null) {
queue.add(matchingChild);
foundLeafNode |= matchingChild.isLeaf();
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index 1f4cdaddcfebf..e5191244df7e1 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -72,7 +72,7 @@ public int hashCode() {
@Override
public String toString(String field) {
- // Does not implements a user-readable toString
+ // Does not implement a user-readable toString
return null;
}
@@ -98,6 +98,7 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
@Override
public boolean isCacheable(LeafReaderContext ctx) {
+ // TODO: Can only cache when segment maxDocs > starTreeDocCount
return false;
}
};
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 9f6c2aa764bff..ebdaf2b924430 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -100,7 +100,6 @@ public void testStarTreeDocValues() throws IOException {
conf.setMergePolicy(newLogMergePolicy());
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
- List<Document> documents = new ArrayList<>();
Random random = new Random();
int totalDocs = 100;
final String SNDV = "sndv";
@@ -112,22 +111,18 @@ public void testStarTreeDocValues() throws IOException {
Document doc = new Document();
if (random.nextBoolean()) {
val = random.nextInt(10) - 5;
- val = val == -1 ? 1 : val; // -1 is encoded as null
doc.add(new SortedNumericDocValuesField(SNDV, val)); // Random long between -5 and 4
}
if (random.nextBoolean()) {
val = random.nextInt(20) - 10;
- val = val == -1 ? 1 : val; // -1 is encoded as null
doc.add(new SortedNumericDocValuesField(DV, val)); // Random long between -10 and 9
}
if (random.nextBoolean()) {
doc.add(new SortedNumericDocValuesField(FIELD_NAME, random.nextInt(50))); // Random long between 0 and 49
}
iw.addDocument(doc);
- documents.add(doc);
}
- iw.forceMerge(1);
if (randomBoolean()) {
iw.forceMerge(1);
}
@@ -159,19 +154,16 @@ public void testStarTreeDocValues() throws IOException {
testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
// numeric-terms query
- for (int cases = 0; cases < 1; cases++) {
+ for (int cases = 0; cases < 100; cases++) {
Map<String, Long> queryMap;
String queryField;
long queryValue;
if (randomBoolean()) {
queryField = SNDV;
- queryValue = random.nextInt(10) - 5;
+ queryValue = random.nextInt(10);
} else {
queryField = DV;
- queryValue = random.nextInt(20) - 10;
- }
- if (queryValue == -1) {
- continue; // -1 is encoded as null
+ queryValue = random.nextInt(20) - 15;
}
defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
queryMap = Map.of(queryField, queryValue);
From 5758d97fe9cb284d23c87fc85cf5fe71c6642f40 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 11:30:14 -0700
Subject: [PATCH 06/35] enable query caching
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/common/util/FeatureFlags.java | 2 +-
.../startree/OriginalOrStarTreeQuery.java | 4 ---
.../search/startree/StarTreeQuery.java | 29 ++++++++++++-------
.../StarTreeDocValuesFormatTests.java | 1 -
.../startree/MetricAggregatorTests.java | 1 -
5 files changed, 20 insertions(+), 17 deletions(-)
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index e663d8429da13..6df68013a8119 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
index 34044d1c44bc7..e673ab91d9e2f 100644
--- a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
@@ -70,10 +70,6 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
}
}
- public Query getOriginalQuery() {
- return originalQuery;
- }
-
public StarTreeQuery getStarTreeQuery() {
return starTreeQuery;
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index e5191244df7e1..6619e46b6c1c2 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -81,16 +81,10 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
return new ConstantScoreWeight(this, boost) {
@Override
public Scorer scorer(LeafReaderContext context) throws IOException {
- SegmentReader reader = Lucene.segmentReader(context.reader());
- CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
- List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
- StarTreeValues starTreeValues = null;
- if (compositeIndexFields != null && !compositeIndexFields.isEmpty()) {
- starTreeValues = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
- } else {
+ StarTreeValues starTreeValues = getStarTreeValues(context);
+ if (starTreeValues == null) {
return null;
}
-
StarTreeFilter filter = new StarTreeFilter(starTreeValues, queryMap);
DocIdSetIterator result = filter.getStarTreeResult();
return new ConstantScoreScorer(this, score(), scoreMode, result);
@@ -98,8 +92,23 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
@Override
public boolean isCacheable(LeafReaderContext ctx) {
- // TODO: Can only cache when segment maxDocs > starTreeDocCount
- return false;
+ try {
+ return getStarTreeValues(ctx).getStarTreeDocumentCount() < ctx.reader().maxDoc();
+ } catch (Exception suppressed) {
+ assert false : "Not able to check cacheable criteria";
+ return false;
+ }
+ }
+
+ private StarTreeValues getStarTreeValues(LeafReaderContext ctx) throws IOException {
+ SegmentReader reader = Lucene.segmentReader(ctx.reader());
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+ List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+ if (compositeIndexFields != null && !compositeIndexFields.isEmpty()) {
+ return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
+ } else {
+ return null;
+ }
}
};
}
diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
index 2aaa815662427..d58535403dd2c 100644
--- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
+++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
@@ -367,7 +367,6 @@ private static XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOExc
return builder.endObject().endObject();
}
-
public static MapperService createMapperService(XContentBuilder builder) throws IOException {
Settings settings = Settings.builder()
.put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index ebdaf2b924430..61c972146b64d 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -52,7 +52,6 @@
import org.junit.Before;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
From 326972186a17d53db9e2e6ea1215d4979a724a69 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 14:12:18 -0700
Subject: [PATCH 07/35] minor refactor
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/index/query/QueryShardContext.java | 13 +++++++------
.../java/org/opensearch/search/SearchService.java | 13 ++-----------
.../search/aggregations/metrics/AvgAggregator.java | 8 +++-----
.../search/aggregations/metrics/MaxAggregator.java | 8 +++-----
.../search/aggregations/metrics/MinAggregator.java | 8 +++-----
.../metrics/NumericMetricsAggregator.java | 10 ++++++++++
.../search/aggregations/metrics/SumAggregator.java | 8 +++-----
.../aggregations/metrics/ValueCountAggregator.java | 8 +++-----
.../startree/MetricAggregatorTests.java | 11 ++++++-----
9 files changed, 40 insertions(+), 47 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index 2539775c3b541..93ec90771f0f9 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -636,17 +636,18 @@ private Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<
}
public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory) {
- String field;
- Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
- .stream()
- .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
+ if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) {
+ String field;
+ Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
+ .stream()
+ .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
- if (aggregatorFactory instanceof MetricAggregatorFactory) {
MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
field = ((MetricAggregatorFactory) aggregatorFactory).getField();
return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
+ } else {
+ return false;
}
- return false;
}
public Index index() {
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 950dad4cd720e..60290b4c78823 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -107,7 +107,6 @@
import org.opensearch.search.aggregations.MultiBucketConsumerService;
import org.opensearch.search.aggregations.SearchContextAggregations;
import org.opensearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
-import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.collapse.CollapseContext;
import org.opensearch.search.deciders.ConcurrentSearchRequestDecider;
@@ -1367,7 +1366,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
// Can be marked false for majority cases for which star-tree cannot be used
// As we increment the cases where star-tree can be used, this can be set back to true
- boolean canUseStarTree = this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
+ boolean canUseStarTree = source.aggregations() != null
+ && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
&& context.mapperService().isCompositeIndexPresent();
SearchShardTarget shardTarget = context.shardTarget();
@@ -1571,11 +1571,6 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source)
throws IOException {
-
- if (source.aggregations() == null) {
- return false;
- }
-
// Current implementation assumes only single star-tree is supported
CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
.getCompositeFieldTypes()
@@ -1592,10 +1587,6 @@ private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryS
}
for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
- if (!(aggregatorFactory instanceof ValuesSourceAggregatorFactory
- && aggregatorFactory.getSubFactories().getFactories().length == 0)) {
- return false;
- }
if (queryShardContext.validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
return false;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index a4338f6eb5812..4382d4a397e98 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -52,8 +52,6 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -101,9 +99,9 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource == null) {
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ if (supportedStarTree != null) {
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index bd511a9f2205d..e33f32352454b 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -57,8 +57,6 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Arrays;
@@ -129,9 +127,9 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ if (supportedStarTree != null) {
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index e2a1bb053f937..ba5c3b7c4f1c9 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -57,8 +57,6 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -128,9 +126,9 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ if (supportedStarTree != null) {
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index 395bf32c7ffe8..cefbbcd5821f8 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -41,6 +41,8 @@
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.sort.SortOrder;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -85,6 +87,14 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
}
return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(lhs), metric(rhs), order == SortOrder.ASC);
}
+
+ public CompositeIndexFieldInfo getSupportedStarTree(LeafReaderContext ctx) {
+ if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+ StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+ return starTreeQuery.getStarTree();
+ }
+ return null;
+ }
}
/**
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 82bdf3030b3f3..a7a409d955d81 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -51,8 +51,6 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -98,9 +96,9 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ if (supportedStarTree != null) {
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index b0cc37d1703fd..b1c5882022ec2 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -50,8 +50,6 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -95,9 +93,9 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource instanceof ValuesSource.Numeric) {
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ if (supportedStarTree != null) {
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx);
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 61c972146b64d..873c3e3e3e5b4 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -109,15 +109,16 @@ public void testStarTreeDocValues() throws IOException {
for (int i = 0; i < totalDocs; i++) {
Document doc = new Document();
if (random.nextBoolean()) {
- val = random.nextInt(10) - 5;
- doc.add(new SortedNumericDocValuesField(SNDV, val)); // Random long between -5 and 4
+ val = random.nextInt(10) - 5; // Random long between -5 and 4
+ doc.add(new SortedNumericDocValuesField(SNDV, val));
}
if (random.nextBoolean()) {
- val = random.nextInt(20) - 10;
- doc.add(new SortedNumericDocValuesField(DV, val)); // Random long between -10 and 9
+ val = random.nextInt(20) - 10; // Random long between -10 and 9
+ doc.add(new SortedNumericDocValuesField(DV, val));
}
if (random.nextBoolean()) {
- doc.add(new SortedNumericDocValuesField(FIELD_NAME, random.nextInt(50))); // Random long between 0 and 49
+ val = random.nextInt(50); // Random long between 0 and 49
+ doc.add(new SortedNumericDocValuesField(FIELD_NAME, val));
}
iw.addDocument(doc);
}
From a25678352c0e00f4979f41bc25b86acc5be6d894 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 18:44:31 -0700
Subject: [PATCH 08/35] use rewrite instead of createweight in
originalorstartreequery
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../aggregations/AggregatorFactories.java | 2 +-
.../metrics/NumericMetricsAggregator.java | 7 +-
.../startree/OriginalOrStarTreeQuery.java | 21 +----
.../search/startree/StarTreeFilter.java | 3 +-
.../search/startree/StarTreeQuery.java | 16 +++-
.../opensearch/search/SearchServiceTests.java | 76 +++++++++++++++++++
.../startree/MetricAggregatorTests.java | 2 +-
.../aggregations/AggregatorTestCase.java | 2 +-
8 files changed, 100 insertions(+), 29 deletions(-)
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
index ab3d78032dc54..720a24da1d9d4 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
@@ -255,7 +255,7 @@ public static Builder builder() {
return new Builder();
}
- AggregatorFactories(AggregatorFactory[] factories) {
+ private AggregatorFactories(AggregatorFactory[] factories) {
this.factories = factories;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index cefbbcd5821f8..21e41616d65c5 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -41,7 +41,6 @@
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
@@ -89,9 +88,9 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
}
public CompositeIndexFieldInfo getSupportedStarTree(LeafReaderContext ctx) {
- if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
- StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
- return starTreeQuery.getStarTree();
+ if (context.query() instanceof StarTreeQuery) {
+ // StarTreeQuery starTreeQuery = ((StarTreeQuery) context.query()).getStarTreeQuery();
+ return ((StarTreeQuery) context.query()).getStarTree();
}
return null;
}
diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
index e673ab91d9e2f..4806265888a2c 100644
--- a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
@@ -11,8 +11,6 @@
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Weight;
import java.io.IOException;
import java.util.Objects;
@@ -27,12 +25,10 @@ public class OriginalOrStarTreeQuery extends Query {
private final StarTreeQuery starTreeQuery;
private final Query originalQuery;
- private boolean starTreeQueryUsed;
public OriginalOrStarTreeQuery(StarTreeQuery starTreeQuery, Query originalQuery) {
this.starTreeQuery = starTreeQuery;
this.originalQuery = originalQuery;
- this.starTreeQueryUsed = false;
}
@Override
@@ -57,20 +53,11 @@ public int hashCode() {
return Objects.hash(classHash(), starTreeQuery, originalQuery, starTreeQuery);
}
- public boolean isStarTreeUsed() {
- return starTreeQueryUsed;
- }
-
- public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
- if (searcher.getIndexReader().hasDeletions() == false) {
- this.starTreeQueryUsed = true;
- return this.starTreeQuery.createWeight(searcher, scoreMode, boost);
- } else {
- return this.originalQuery.createWeight(searcher, scoreMode, boost);
+ @Override
+ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+ if (indexSearcher.getIndexReader().hasDeletions()) {
+ return originalQuery;
}
- }
-
- public StarTreeQuery getStarTreeQuery() {
return starTreeQuery;
}
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index b2838ad788da6..4ff0e9190ca35 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -36,8 +36,9 @@
* Filter operator for star tree data structure.
*
* @opensearch.experimental
+ * @opensearch.internal
*/
-public class StarTreeFilter {
+class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
// private final StarTreeNode starTreeRoot;
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index 6619e46b6c1c2..f292ea2b31a35 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -40,13 +40,13 @@ public class StarTreeQuery extends Query {
* Star tree field info
* This is used to get the star tree data structure
*/
- CompositeIndexFieldInfo starTree;
+ private final CompositeIndexFieldInfo starTree;
/**
* Map of field name to a value to be queried for that field
* This is used to filter the data based on the query
*/
- Map<String, Long> queryMap;
+ private final Map<String, Long> queryMap;
public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
this.starTree = starTree;
@@ -72,8 +72,16 @@ public int hashCode() {
@Override
public String toString(String field) {
- // Does not implement a user-readable toString
- return null;
+ StringBuilder sb = new StringBuilder();
+ sb.append(getClass().getSimpleName());
+ sb.append("(");
+ sb.append(this.starTree);
+ if (queryMap != null) {
+ sb.append(", ");
+ sb.append(queryMap);
+ sb.append(")");
+ }
+ return sb.toString();
}
@Override
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index 514e99a126267..301f6b0c5aa44 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -35,10 +35,12 @@
import org.apache.lucene.index.FilterDirectoryReader;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.store.AlreadyClosedException;
import org.opensearch.OpenSearchException;
import org.opensearch.action.OriginalIndices;
+import org.opensearch.action.admin.indices.create.CreateIndexRequestBuilder;
import org.opensearch.action.index.IndexResponse;
import org.opensearch.action.search.ClearScrollRequest;
import org.opensearch.action.search.DeletePitResponse;
@@ -54,16 +56,20 @@
import org.opensearch.action.support.IndicesOptions;
import org.opensearch.action.support.PlainActionFuture;
import org.opensearch.action.support.WriteRequest;
+import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.UUIDs;
import org.opensearch.common.settings.Setting;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.settings.SettingsException;
import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.FeatureFlags;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.unit.ByteSizeUnit;
+import org.opensearch.core.common.unit.ByteSizeValue;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
@@ -73,6 +79,9 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
+import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.compositeindex.CompositeIndexSettings;
+import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.mapper.DerivedFieldType;
import org.opensearch.index.query.AbstractQueryBuilder;
@@ -113,6 +122,7 @@
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.MinAndMax;
import org.opensearch.search.sort.SortOrder;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
import org.opensearch.search.suggest.SuggestBuilder;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
import org.opensearch.threadpool.ThreadPool;
@@ -2259,4 +2269,70 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
primarySort.order(SortOrder.DESC);
assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true);
}
+
+ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
+ FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+
+ Settings enableStarTree = Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), true).build();
+ client().admin().cluster().prepareUpdateSettings().setPersistentSettings(enableStarTree).execute();
+
+ Settings settings = Settings.builder()
+ .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+ .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+ .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
+ .put(IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING.getKey(), new ByteSizeValue(512, ByteSizeUnit.MB))
+ .build();
+ CreateIndexRequestBuilder builder = client().admin()
+ .indices()
+ .prepareCreate("test")
+ .setSettings(settings)
+ .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
+ createIndex("test", builder);
+ IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+ IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
+ IndexShard indexShard = indexService.getShard(0);
+
+ SearchService searchService = getInstanceFromNode(SearchService.class);
+
+ // Case 1: No query or aggregations, should not use star tree
+ SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+ ShardSearchRequest request = new ShardSearchRequest(
+ OriginalIndices.NONE,
+ new SearchRequest().allowPartialSearchResults(true),
+ indexShard.shardId(),
+ 1,
+ new AliasFilter(null, Strings.EMPTY_ARRAY),
+ 1.0f,
+ -1,
+ null,
+ null
+ );
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
+ SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
+ assertFalse(context.query() instanceof OriginalOrStarTreeQuery);
+ }
+
+ // Case 2: Query present but no aggregations, should not use star tree
+ sourceBuilder.query(new MatchAllQueryBuilder());
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
+ SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
+ assertThat(context.query(), instanceOf(MatchAllDocsQuery.class));
+ }
+
+ // Case 3: Query and aggregations present, should use star tree if possible
+ sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
+ SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
+ if (context.query() instanceof OriginalOrStarTreeQuery) {
+ // Star tree query was set
+ assertThat(context.query(), instanceOf(OriginalOrStarTreeQuery.class));
+ } else {
+ // Star tree query was not set (e.g., no star tree index present)
+ assertThat(context.query(), instanceOf(MatchAllQueryBuilder.class));
+ }
+ }
+ }
+
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 873c3e3e3e5b4..500eb47a65ae0 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -185,7 +185,7 @@ public void testStarTreeDocValues() throws IOException {
directory.close();
}
- private <T extends AggregationBuilder, V extends InternalAggregation, R extends Number> void testCase(
+ private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
IndexSearcher searcher,
Query defaultQuery,
StarTreeQuery starTreeQuery,
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 6f9b6330ac95e..d4a37bfd87bd3 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -661,6 +661,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
boolean hasNested,
MappedFieldType... fieldTypes
) throws IOException {
+ query = query.rewrite(searcher);
final IndexReaderContext ctx = searcher.getTopReaderContext();
final PipelineTree pipelines = builder.buildPipelineTree();
List<InternalAggregation> aggs = new ArrayList<>();
@@ -692,7 +693,6 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
a.preCollection();
Weight weight = subSearcher.createWeight(query, ScoreMode.COMPLETE, 1f);
-// assertTrue(weight.getQuery() instanceof StarTreeQuery);
subSearcher.search(weight, a);
a.postCollection();
aggs.add(a.buildTopLevel());
From bfe347fb5a706b7c17076df944d7589dab3e1c3e Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 22:29:52 -0700
Subject: [PATCH 09/35] fix request parsing tests
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../index/query/QueryShardContext.java | 3 +-
.../org/opensearch/search/SearchService.java | 1 +
.../metrics/NumericMetricsAggregator.java | 4 -
.../opensearch/search/SearchServiceTests.java | 88 ++++++++++++++-----
.../startree/MetricAggregatorTests.java | 4 +-
5 files changed, 72 insertions(+), 28 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index 93ec90771f0f9..9f7aeb7732d8d 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -35,6 +35,7 @@
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.join.BitSetProducer;
import org.apache.lucene.search.similarities.Similarity;
@@ -599,7 +600,7 @@ public ParsedQuery toStarTreeQuery(
Query query
) {
Map<String, Long> queryMap;
- if (queryBuilder == null) {
+ if (queryBuilder == null || query instanceof MatchAllDocsQuery) {
queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 60290b4c78823..38e751f4a0b08 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -1367,6 +1367,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
// Can be marked false for majority cases for which star-tree cannot be used
// As we increment the cases where star-tree can be used, this can be set back to true
boolean canUseStarTree = source.aggregations() != null
+ && includeAggregations
&& this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
&& context.mapperService().isCompositeIndexPresent();
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index 21e41616d65c5..4e4172d910f14 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -45,7 +45,6 @@
import java.io.IOException;
import java.util.Map;
-import java.util.concurrent.atomic.AtomicReference;
/**
* Base class to aggregate all docs into a single numeric metric value.
@@ -89,7 +88,6 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
public CompositeIndexFieldInfo getSupportedStarTree(LeafReaderContext ctx) {
if (context.query() instanceof StarTreeQuery) {
- // StarTreeQuery starTreeQuery = ((StarTreeQuery) context.query()).getStarTreeQuery();
return ((StarTreeQuery) context.query()).getStarTree();
}
return null;
@@ -131,8 +129,6 @@ protected StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndex
}
CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
StarTreeValues values = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
- final AtomicReference<StarTreeValues> aggrVal = new AtomicReference<>(null);
-
return values;
}
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index 301f6b0c5aa44..ddeab0ef1eb59 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -68,8 +68,6 @@
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
-import org.opensearch.core.common.unit.ByteSizeUnit;
-import org.opensearch.core.common.unit.ByteSizeValue;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
@@ -122,7 +120,7 @@
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.MinAndMax;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.search.suggest.SuggestBuilder;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
import org.opensearch.threadpool.ThreadPool;
@@ -2272,15 +2270,16 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
-
- Settings enableStarTree = Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), true).build();
- client().admin().cluster().prepareUpdateSettings().setPersistentSettings(enableStarTree).execute();
+ client().admin()
+ .cluster()
+ .prepareUpdateSettings()
+ .setTransientSettings(Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), true).build())
+ .execute();
Settings settings = Settings.builder()
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
.put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
.put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
- .put(IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING.getKey(), new ByteSizeValue(512, ByteSizeUnit.MB))
.build();
CreateIndexRequestBuilder builder = client().admin()
.indices()
@@ -2308,31 +2307,76 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
null
);
try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
- SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
- assertFalse(context.query() instanceof OriginalOrStarTreeQuery);
+ SearchContext context = searchService.createContext(reader, request, null, false);
+ assertFalse(context.query() instanceof StarTreeQuery);
+ searchService.doStop();
}
- // Case 2: Query present but no aggregations, should not use star tree
+ // Case 2: MatchAllQuery present but no aggregations, should not use star tree
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
sourceBuilder.query(new MatchAllQueryBuilder());
request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
assertThat(context.query(), instanceOf(MatchAllDocsQuery.class));
+ searchService.doStop();
}
- // Case 3: Query and aggregations present, should use star tree if possible
+ // Case 3: MatchAllQuery and aggregations present, should use star tree if possible
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
+ sourceBuilder.query(new MatchAllQueryBuilder());
sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
- SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
- if (context.query() instanceof OriginalOrStarTreeQuery) {
- // Star tree query was set
- assertThat(context.query(), instanceOf(OriginalOrStarTreeQuery.class));
- } else {
- // Star tree query was not set (e.g., no star tree index present)
- assertThat(context.query(), instanceOf(MatchAllQueryBuilder.class));
- }
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(StarTreeQuery.class));
+ searchService.doStop();
}
- }
+ // Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
+ sourceBuilder.query(new MatchAllQueryBuilder());
+ sourceBuilder.aggregation(AggregationBuilders.max("test").field("field")).trackTotalHitsUpTo(1000);
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(MatchAllDocsQuery.class));
+ }
+ searchService.doStop();
+
+ // Case 5: TermQuery and aggregations present, should use star tree
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
+ sourceBuilder.query(new TermQueryBuilder("sndv", 1));
+ sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(StarTreeQuery.class));
+ }
+ searchService.doStop();
+
+ // Case 6: No query, metric aggregations present, should use star tree
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
+ sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(StarTreeQuery.class));
+ }
+ searchService.doStop();
+
+ searchService.doClose();
+ client().admin()
+ .cluster()
+ .prepareUpdateSettings()
+ .setTransientSettings(
+ Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), (String) null).build()
+ )
+ .execute();
+ }
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 500eb47a65ae0..41eba6b548b45 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -8,6 +8,8 @@
package org.opensearch.search.aggregations.startree;
+import com.carrotsearch.randomizedtesting.RandomizedTest;
+
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.codecs.Codec;
@@ -99,7 +101,7 @@ public void testStarTreeDocValues() throws IOException {
conf.setMergePolicy(newLogMergePolicy());
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
- Random random = new Random();
+ Random random = RandomizedTest.getRandom();
int totalDocs = 100;
final String SNDV = "sndv";
final String DV = "dv";
From 18aaacef26155431be0d0423e14eff800273e17a Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 3 Sep 2024 23:02:50 -0700
Subject: [PATCH 10/35] minor refactoring and test:framework spotless fix
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/search/aggregations/metrics/AvgAggregator.java | 2 +-
.../opensearch/search/aggregations/metrics/MaxAggregator.java | 2 +-
.../opensearch/search/aggregations/metrics/MinAggregator.java | 2 +-
.../search/aggregations/metrics/NumericMetricsAggregator.java | 2 +-
.../opensearch/search/aggregations/metrics/SumAggregator.java | 2 +-
.../search/aggregations/metrics/ValueCountAggregator.java | 2 +-
.../org/opensearch/search/aggregations/AggregatorTestCase.java | 3 +--
7 files changed, 7 insertions(+), 8 deletions(-)
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 4382d4a397e98..12742bd681953 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -99,7 +99,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource == null) {
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index e33f32352454b..1a37682d7b97c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -127,7 +127,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index ba5c3b7c4f1c9..9a7906f75624b 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -126,7 +126,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index 4e4172d910f14..e5509b4d849bc 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -86,7 +86,7 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(lhs), metric(rhs), order == SortOrder.ASC);
}
- public CompositeIndexFieldInfo getSupportedStarTree(LeafReaderContext ctx) {
+ public CompositeIndexFieldInfo getSupportedStarTree() {
if (context.query() instanceof StarTreeQuery) {
return ((StarTreeQuery) context.query()).getStarTree();
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index a7a409d955d81..6c77a8381ce11 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -96,7 +96,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index b1c5882022ec2..bdf8d09c704a7 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -93,7 +93,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource instanceof ValuesSource.Numeric) {
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree(ctx);
+ CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index d4a37bfd87bd3..63288602abb7e 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -37,7 +37,6 @@
import org.apache.lucene.document.Field;
import org.apache.lucene.document.InetAddressPoint;
import org.apache.lucene.document.LatLonDocValuesField;
-import org.apache.lucene.document.LongField;
import org.apache.lucene.document.SortedNumericDocValuesField;
import org.apache.lucene.document.SortedSetDocValuesField;
import org.apache.lucene.document.StoredField;
@@ -142,7 +141,6 @@
import org.opensearch.search.internal.ContextIndexSearcher;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.lookup.SearchLookup;
-import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.test.InternalAggregationTestCase;
import org.opensearch.test.OpenSearchTestCase;
import org.junit.After;
@@ -652,6 +650,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer);
return internalAgg;
}
+
protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduceStarTree(
IndexSettings indexSettings,
IndexSearcher searcher,
From 2857939053b9df1dc435f7e9a094c60b1238cab5 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 00:43:47 -0700
Subject: [PATCH 11/35] Refactoring star tree query utils in a utility class
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 143 ++++++++++++++++++
.../index/query/QueryShardContext.java | 69 ---------
.../org/opensearch/search/SearchService.java | 52 +------
.../StarTreeDocValuesFormatTests.java | 5 +-
4 files changed, 153 insertions(+), 116 deletions(-)
create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
new file mode 100644
index 0000000000000..867325b27e422
--- /dev/null
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -0,0 +1,143 @@
+/*
+ * 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.index.compositeindex.datacube.startree.utils;
+
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.Dimension;
+import org.opensearch.index.compositeindex.datacube.Metric;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
+import org.opensearch.index.mapper.StarTreeMapper;
+import org.opensearch.index.query.MatchAllQueryBuilder;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * Helper class for building star-tree query
+ *
+ * @opensearch.internal
+ * @opensearch.experimental
+ */
+public class StarTreeQueryHelper {
+
+ /**
+ * Checks if the search context can be supported by star-tree
+ */
+ public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) {
+ boolean canUseStarTree = context.aggregations() != null
+ && context.mapperService().isCompositeIndexPresent()
+ && context.parsedPostFilter() == null
+ && context.innerHits().getInnerHits().isEmpty()
+ && context.sort() == null
+ && (!trackTotalHits || context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_DISABLED)
+ && context.trackScores() == false
+ && context.minimumScore() == null
+ && context.terminateAfter() == 0;
+ return canUseStarTree;
+ }
+
+ /**
+ * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
+ * Returns null if the query cannot be supported.
+ */
+ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext context, SearchSourceBuilder source) throws IOException {
+ // Current implementation assumes only single star-tree is supported
+ CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
+ .getCompositeFieldTypes()
+ .iterator()
+ .next();
+ CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo(
+ compositeMappedFieldType.name(),
+ compositeMappedFieldType.getCompositeIndexType()
+ );
+
+ StarTreeQuery starTreeQuery = StarTreeQueryHelper.toStarTreeQuery(starTree, compositeMappedFieldType, source.query());
+ if (starTreeQuery == null) {
+ return null;
+ }
+
+ for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
+ if (validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
+ return null;
+ }
+ }
+
+ return new OriginalOrStarTreeQuery(starTreeQuery, context.query());
+ }
+
+ private static StarTreeQuery toStarTreeQuery(
+ CompositeIndexFieldInfo starTree,
+ CompositeDataCubeFieldType compositeIndexFieldInfo,
+ QueryBuilder queryBuilder
+ ) {
+ Map<String, Long> queryMap;
+ if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
+ queryMap = null;
+ } else if (queryBuilder instanceof TermQueryBuilder) {
+ List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
+ .stream()
+ .map(Dimension::getField)
+ .collect(Collectors.toList());
+ queryMap = getStarTreePredicates(queryBuilder, supportedDimensions);
+ } else {
+ return null;
+ }
+
+ return new StarTreeQuery(starTree, queryMap);
+ }
+
+ /**
+ * Parse query body to star-tree predicates
+ * @param queryBuilder
+ * @return predicates to match
+ */
+ private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
+ TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
+ String field = tq.fieldName();
+ if (!supportedDimensions.contains(field)) {
+ throw new IllegalArgumentException("unsupported field in star-tree");
+ }
+ long inputQueryVal = Long.parseLong(tq.value().toString());
+
+ // Create a map with the field and the value
+ Map<String, Long> predicateMap = new HashMap<>();
+ predicateMap.put(field, inputQueryVal);
+ return predicateMap;
+ }
+
+ private static boolean validateStarTreeMetricSuport(
+ CompositeDataCubeFieldType compositeIndexFieldInfo,
+ AggregatorFactory aggregatorFactory
+ ) {
+ if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) {
+ String field;
+ Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
+ .stream()
+ .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
+
+ MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
+ field = ((MetricAggregatorFactory) aggregatorFactory).getField();
+ return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
+ } else {
+ return false;
+ }
+ }
+}
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index 9f7aeb7732d8d..bccead2b029d0 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -35,7 +35,6 @@
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.join.BitSetProducer;
import org.apache.lucene.search.similarities.Similarity;
@@ -57,12 +56,7 @@
import org.opensearch.index.IndexSortConfig;
import org.opensearch.index.analysis.IndexAnalyzers;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.compositeindex.datacube.Dimension;
-import org.opensearch.index.compositeindex.datacube.Metric;
-import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.fielddata.IndexFieldData;
-import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.ContentPath;
import org.opensearch.index.mapper.DerivedFieldResolver;
import org.opensearch.index.mapper.DerivedFieldResolverFactory;
@@ -79,13 +73,9 @@
import org.opensearch.script.ScriptContext;
import org.opensearch.script.ScriptFactory;
import org.opensearch.script.ScriptService;
-import org.opensearch.search.aggregations.AggregatorFactory;
-import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
import org.opensearch.search.aggregations.support.AggregationUsageService;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.lookup.SearchLookup;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.transport.RemoteClusterAware;
import java.io.IOException;
@@ -99,7 +89,6 @@
import java.util.function.LongSupplier;
import java.util.function.Predicate;
import java.util.function.Supplier;
-import java.util.stream.Collectors;
import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
@@ -593,64 +582,6 @@ private ParsedQuery toQuery(QueryBuilder queryBuilder, CheckedFunction<QueryBuil
}
}
- public ParsedQuery toStarTreeQuery(
- CompositeIndexFieldInfo starTree,
- CompositeDataCubeFieldType compositeIndexFieldInfo,
- QueryBuilder queryBuilder,
- Query query
- ) {
- Map<String, Long> queryMap;
- if (queryBuilder == null || query instanceof MatchAllDocsQuery) {
- queryMap = null;
- } else if (queryBuilder instanceof TermQueryBuilder) {
- List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
- .stream()
- .map(Dimension::getField)
- .collect(Collectors.toList());
- queryMap = getStarTreePredicates(queryBuilder, supportedDimensions);
- } else {
- return null;
- }
-
- StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, queryMap);
- OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, query);
- return new ParsedQuery(originalOrStarTreeQuery);
- }
-
- /**
- * Parse query body to star-tree predicates
- * @param queryBuilder
- * @return predicates to match
- */
- private Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
- TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
- String field = tq.fieldName();
- if (!supportedDimensions.contains(field)) {
- throw new IllegalArgumentException("unsupported field in star-tree");
- }
- long inputQueryVal = Long.parseLong(tq.value().toString());
-
- // Create a map with the field and the value
- Map<String, Long> predicateMap = new HashMap<>();
- predicateMap.put(field, inputQueryVal);
- return predicateMap;
- }
-
- public boolean validateStarTreeMetricSuport(CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory) {
- if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) {
- String field;
- Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
- .stream()
- .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
-
- MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
- field = ((MetricAggregatorFactory) aggregatorFactory).getField();
- return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
- } else {
- return false;
- }
- }
-
public Index index() {
return indexSettings.getIndex();
}
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 38e751f4a0b08..a1cb51f098a2a 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -77,12 +77,10 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.engine.Engine;
-import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.DerivedFieldResolver;
import org.opensearch.index.mapper.DerivedFieldResolverFactory;
-import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.InnerHitContextBuilder;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.MatchNoneQueryBuilder;
@@ -101,7 +99,6 @@
import org.opensearch.script.ScriptService;
import org.opensearch.search.aggregations.AggregationInitializationException;
import org.opensearch.search.aggregations.AggregatorFactories;
-import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.InternalAggregation;
import org.opensearch.search.aggregations.InternalAggregation.ReduceContext;
import org.opensearch.search.aggregations.MultiBucketConsumerService;
@@ -142,6 +139,7 @@
import org.opensearch.search.sort.SortAndFormats;
import org.opensearch.search.sort.SortBuilder;
import org.opensearch.search.sort.SortOrder;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
import org.opensearch.search.suggest.Suggest;
import org.opensearch.search.suggest.completion.CompletionSuggestion;
import org.opensearch.tasks.TaskResourceTrackingService;
@@ -1364,12 +1362,6 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.evaluateRequestShouldUseConcurrentSearch();
return;
}
- // Can be marked false for majority cases for which star-tree cannot be used
- // As we increment the cases where star-tree can be used, this can be set back to true
- boolean canUseStarTree = source.aggregations() != null
- && includeAggregations
- && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
- && context.mapperService().isCompositeIndexPresent();
SearchShardTarget shardTarget = context.shardTarget();
QueryShardContext queryShardContext = context.getQueryShardContext();
@@ -1381,12 +1373,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.parsedQuery(queryShardContext.toQuery(source.query()));
}
if (source.postFilter() != null) {
- canUseStarTree = false;
InnerHitContextBuilder.extractInnerHits(source.postFilter(), innerHitBuilders);
context.parsedPostFilter(queryShardContext.toQuery(source.postFilter()));
}
if (!innerHitBuilders.isEmpty()) {
- canUseStarTree = false;
for (Map.Entry<String, InnerHitContextBuilder> entry : innerHitBuilders.entrySet()) {
try {
entry.getValue().build(context, context.innerHits());
@@ -1396,7 +1386,6 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
}
if (source.sorts() != null) {
- canUseStarTree = false;
try {
Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(source.sorts(), context.getQueryShardContext());
optionalSort.ifPresent(context::sort);
@@ -1413,12 +1402,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
if (source.trackTotalHitsUpTo() != null) {
context.trackTotalHitsUpTo(source.trackTotalHitsUpTo());
- if (source.trackTotalHitsUpTo() != TRACK_TOTAL_HITS_DISABLED) {
- canUseStarTree = false;
- }
}
if (source.minScore() != null) {
- canUseStarTree = false;
context.minimumScore(source.minScore());
}
if (source.timeout() != null) {
@@ -1559,9 +1544,12 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch()));
}
- if (canUseStarTree) {
+ if (this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
+ && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) {
try {
- if (setStarTreeQuery(context, queryShardContext, source)) {
+ OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source);
+ if (parsedQuery != null) {
+ context.parsedQuery(new ParsedQuery(parsedQuery));
logger.debug("can use star tree");
} else {
logger.debug("cannot use star tree");
@@ -1570,32 +1558,6 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
}
}
- private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source)
- throws IOException {
- // Current implementation assumes only single star-tree is supported
- CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
- .getCompositeFieldTypes()
- .iterator()
- .next();
- CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo(
- compositeMappedFieldType.name(),
- compositeMappedFieldType.getCompositeIndexType()
- );
-
- ParsedQuery newParsedQuery = queryShardContext.toStarTreeQuery(starTree, compositeMappedFieldType, source.query(), context.query());
- if (newParsedQuery == null) {
- return false;
- }
-
- for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
- if (queryShardContext.validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
- return false;
- }
- }
- context.parsedQuery(newParsedQuery);
- return true;
- }
-
/**
* Shortcut ids to load, we load only "from" and up to "size". The phase controller
* handles this as well since the result is always size * shards for Q_T_F
diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
index d58535403dd2c..628458ccb88f8 100644
--- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
+++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
@@ -109,7 +109,7 @@ protected Codec getCodec() {
final Logger testLogger = LogManager.getLogger(StarTreeDocValuesFormatTests.class);
try {
- createMapperService(getExpandedMapping());
+ mapperService = createMapperService(getExpandedMapping());
} catch (IOException e) {
throw new RuntimeException(e);
}
@@ -117,7 +117,8 @@ protected Codec getCodec() {
return codec;
}
- public void testStarTreeDocValues() throws IOException {
+ // TODO: Awaiting a fix in indexing - disabling test for meantime
+ private void testStarTreeDocValues() throws IOException {
Directory directory = newDirectory();
IndexWriterConfig conf = newIndexWriterConfig(null);
conf.setMergePolicy(newLogMergePolicy());
From 6796be14377b5af9228c9aa603603501cc7b8343 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 13:06:07 -0700
Subject: [PATCH 12/35] refactoring to utils
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 27 ++++++++++++++++++-
.../aggregations/metrics/AvgAggregator.java | 10 ++++---
.../aggregations/metrics/MaxAggregator.java | 10 ++++---
.../aggregations/metrics/MinAggregator.java | 10 ++++---
.../metrics/NumericMetricsAggregator.java | 24 -----------------
.../aggregations/metrics/SumAggregator.java | 10 ++++---
.../metrics/ValueCountAggregator.java | 6 ++++-
.../search/startree/StarTreeFilter.java | 16 +++--------
.../search/startree/StarTreeQuery.java | 7 +----
9 files changed, 64 insertions(+), 56 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 867325b27e422..4c76551efe8d3 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -8,10 +8,15 @@
package org.opensearch.index.compositeindex.datacube.startree.utils;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.Metric;
import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.MatchAllQueryBuilder;
@@ -43,6 +48,7 @@ public class StarTreeQueryHelper {
*/
public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) {
boolean canUseStarTree = context.aggregations() != null
+// && context.size() == 0
&& context.mapperService().isCompositeIndexPresent()
&& context.parsedPostFilter() == null
&& context.innerHits().getInnerHits().isEmpty()
@@ -97,6 +103,9 @@ private static StarTreeQuery toStarTreeQuery(
.map(Dimension::getField)
.collect(Collectors.toList());
queryMap = getStarTreePredicates(queryBuilder, supportedDimensions);
+ if (queryMap == null) {
+ return null;
+ }
} else {
return null;
}
@@ -113,7 +122,7 @@ private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder
TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
String field = tq.fieldName();
if (!supportedDimensions.contains(field)) {
- throw new IllegalArgumentException("unsupported field in star-tree");
+ return null;
}
long inputQueryVal = Long.parseLong(tq.value().toString());
@@ -140,4 +149,20 @@ private static boolean validateStarTreeMetricSuport(
return false;
}
}
+
+ public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context) {
+ if (context.query() instanceof StarTreeQuery) {
+ return ((StarTreeQuery) context.query()).getStarTree();
+ }
+ return null;
+ }
+
+ public static StarTreeValues getStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException {
+ SegmentReader reader = Lucene.segmentReader(context.reader());
+ if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) {
+ return null;
+ }
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+ return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 12742bd681953..5a61c559e6e56 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -34,13 +34,13 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.common.util.LongArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -56,6 +56,9 @@
import java.io.IOException;
import java.util.Map;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
+
/**
* Aggregate all docs into an average
*
@@ -99,7 +102,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource == null) {
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
+ CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
@@ -153,6 +156,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.SUM.getTypeName()
);
+ assert starTreeValues != null;
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
@@ -180,7 +184,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valueCount; i++) {
- double value = NumericUtils.sortableLongToDouble(values.nextValue());
+ double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 1a37682d7b97c..98cd48c34ef7e 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -38,12 +38,12 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -63,6 +63,9 @@
import java.util.Map;
import java.util.function.Function;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
+
/**
* Aggregate all docs into a max value
*
@@ -127,7 +130,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
+ CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
@@ -168,6 +171,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.MAX.getTypeName()
);
+ assert starTreeValues != null;
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
@@ -182,7 +186,7 @@ public void collect(int doc, long bucket) throws IOException {
maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = NumericUtils.sortableLongToDouble(values.nextValue());
+ final double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
double max = maxes.get(bucket);
max = Math.max(max, value);
maxes.set(bucket, max);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 9a7906f75624b..60173c2121f26 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -38,12 +38,12 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -62,6 +62,9 @@
import java.util.Map;
import java.util.function.Function;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
+
/**
* Aggregate all docs into a min value
*
@@ -126,7 +129,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
+ CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
@@ -165,6 +168,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.MIN.getTypeName()
);
+ assert starTreeValues != null;
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
@@ -179,7 +183,7 @@ public void collect(int doc, long bucket) throws IOException {
mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = NumericUtils.sortableLongToDouble(values.nextValue());
+ final double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
double min = mins.get(bucket);
min = Math.min(min, value);
mins.set(bucket, min);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
index e5509b4d849bc..f90e5a092385f 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java
@@ -31,17 +31,10 @@
package org.opensearch.search.aggregations.metrics;
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SegmentReader;
-import org.opensearch.common.lucene.Lucene;
import org.opensearch.common.util.Comparators;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.codec.composite.CompositeIndexReader;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.search.aggregations.Aggregator;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.StarTreeQuery;
import java.io.IOException;
import java.util.Map;
@@ -85,13 +78,6 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
}
return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(lhs), metric(rhs), order == SortOrder.ASC);
}
-
- public CompositeIndexFieldInfo getSupportedStarTree() {
- if (context.query() instanceof StarTreeQuery) {
- return ((StarTreeQuery) context.query()).getStarTree();
- }
- return null;
- }
}
/**
@@ -121,14 +107,4 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(key, lhs), metric(key, rhs), order == SortOrder.ASC);
}
}
-
- protected StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
- SegmentReader reader = Lucene.segmentReader(ctx.reader());
- if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) {
- return null;
- }
- CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
- StarTreeValues values = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
- return values;
- }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 6c77a8381ce11..5fd277c3e1c61 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -34,12 +34,12 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -55,6 +55,9 @@
import java.io.IOException;
import java.util.Map;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
+
/**
* Aggregate all docs into a single sum value
*
@@ -96,7 +99,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
return LeafBucketCollector.NO_OP_COLLECTOR;
}
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
+ CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
@@ -142,6 +145,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.SUM.getTypeName()
);
+ assert starTreeValues != null;
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
@@ -160,7 +164,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valuesCount; i++) {
- double value = NumericUtils.sortableLongToDouble(values.nextValue());
+ double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index bdf8d09c704a7..0d9a145b2b8ad 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -54,6 +54,9 @@
import java.io.IOException;
import java.util.Map;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
+
/**
* A field data based aggregator that counts the number of values a specific field has within the aggregation context.
* <p>
@@ -93,7 +96,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (valuesSource instanceof ValuesSource.Numeric) {
- CompositeIndexFieldInfo supportedStarTree = this.getSupportedStarTree();
+ CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
@@ -146,6 +149,7 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf
fieldName,
MetricStat.VALUE_COUNT.getTypeName()
);
+ assert starTreeValues != null;
SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
final BigArrays bigArrays = context.bigArrays();
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 4ff0e9190ca35..7b48ecd4561c7 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -41,23 +41,13 @@
class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
- // private final StarTreeNode starTreeRoot;
- Map<String, Long> queryMap;
- DocIdSetBuilder docsWithField;
- DocIdSetBuilder.BulkAdder adder;
- // Map<String, DocIdSetIterator> dimValueMap;
- StarTreeValues starTreeValues;
- List<Dimension> dimensions;
+ private final Map<String, Long> queryMap;
+ private final StarTreeValues starTreeValues;
public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> predicateEvaluators) {
// This filter operator does not support AND/OR/NOT operations as of now.
starTreeValues = starTreeAggrStructure;
- // starTreeRoot = starTreeAggrStructure.getRoot();
- // dimValueMap = starTreeAggrStructure.getDimensionDocValuesIteratorMap();
queryMap = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
-
- // TODO : this should be the maximum number of doc values
- docsWithField = new DocIdSetBuilder(Integer.MAX_VALUE);
}
/**
@@ -113,6 +103,8 @@ public DocIdSetIterator getStarTreeResult() throws IOException {
* predicate dimensions that are not matched.
*/
private StarTreeResult traverseStarTree() throws IOException {
+ DocIdSetBuilder docsWithField = new DocIdSetBuilder(this.starTreeValues.getStarTreeDocumentCount());
+ DocIdSetBuilder.BulkAdder adder;
Set<String> globalRemainingPredicateColumns = null;
StarTreeNode starTree = starTreeValues.getRoot();
List<String> dimensionNames = starTreeValues.getStarTreeField()
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
index f292ea2b31a35..e8354b0415ac1 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
@@ -100,12 +100,7 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
@Override
public boolean isCacheable(LeafReaderContext ctx) {
- try {
- return getStarTreeValues(ctx).getStarTreeDocumentCount() < ctx.reader().maxDoc();
- } catch (Exception suppressed) {
- assert false : "Not able to check cacheable criteria";
- return false;
- }
+ return false;
}
private StarTreeValues getStarTreeValues(LeafReaderContext ctx) throws IOException {
From 34aa57ac65e5e71653e92884df5c97f63b12da5c Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 17:15:37 -0700
Subject: [PATCH 13/35] rebasing with main
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/search/aggregations/metrics/AvgAggregator.java | 4 ++--
.../opensearch/search/aggregations/metrics/MaxAggregator.java | 4 ++--
.../opensearch/search/aggregations/metrics/MinAggregator.java | 4 ++--
.../opensearch/search/aggregations/metrics/SumAggregator.java | 4 ++--
.../datacube/startree/StarTreeDocValuesFormatTests.java | 3 +--
5 files changed, 9 insertions(+), 10 deletions(-)
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 5a61c559e6e56..d028f38be1b61 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -34,13 +34,13 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.common.util.LongArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -184,7 +184,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valueCount; i++) {
- double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
+ double value = NumericUtils.sortableLongToDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 98cd48c34ef7e..5f0c5e242a453 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -38,12 +38,12 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -186,7 +186,7 @@ public void collect(int doc, long bucket) throws IOException {
maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
+ final double value = NumericUtils.sortableLongToDouble(values.nextValue());
double max = maxes.get(bucket);
max = Math.max(max, value);
maxes.set(bucket, max);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 60173c2121f26..81246a192555c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -38,12 +38,12 @@
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.NumericDoubleValues;
@@ -183,7 +183,7 @@ public void collect(int doc, long bucket) throws IOException {
mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
}
if (values.advanceExact(doc)) {
- final double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
+ final double value = NumericUtils.sortableLongToDouble(values.nextValue());
double min = mins.get(bucket);
min = Math.min(min, value);
mins.set(bucket, min);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 5fd277c3e1c61..c205c9abb7bee 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -34,12 +34,12 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericTypeConverters;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -164,7 +164,7 @@ public void collect(int doc, long bucket) throws IOException {
kahanSummation.reset(sum, compensation);
for (int i = 0; i < valuesCount; i++) {
- double value = StarTreeNumericTypeConverters.sortableLongtoDouble(values.nextValue());
+ double value = NumericUtils.sortableLongToDouble(values.nextValue());
kahanSummation.add(value);
}
diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
index 628458ccb88f8..42ee5a5336dd1 100644
--- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
+++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
@@ -117,8 +117,7 @@ protected Codec getCodec() {
return codec;
}
- // TODO: Awaiting a fix in indexing - disabling test for meantime
- private void testStarTreeDocValues() throws IOException {
+ public void testStarTreeDocValues() throws IOException {
Directory directory = newDirectory();
IndexWriterConfig conf = newIndexWriterConfig(null);
conf.setMergePolicy(newLogMergePolicy());
From a53e11f12405c90ed954b6428088fc07b733f28d Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 19:31:24 -0700
Subject: [PATCH 14/35] minor spotless
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/startree/utils/StarTreeQueryHelper.java | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 4c76551efe8d3..351f38afff7fd 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -48,7 +48,7 @@ public class StarTreeQueryHelper {
*/
public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) {
boolean canUseStarTree = context.aggregations() != null
-// && context.size() == 0
+ && context.size() == 0
&& context.mapperService().isCompositeIndexPresent()
&& context.parsedPostFilter() == null
&& context.innerHits().getInnerHits().isEmpty()
From 19f11c1510ef5fc1799373cb402adc7007b3682b Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 22:17:18 -0700
Subject: [PATCH 15/35] fix search service tests
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/search/SearchServiceTests.java | 27 ++++++++++++++-----
1 file changed, 20 insertions(+), 7 deletions(-)
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index ddeab0ef1eb59..029e3fa302fec 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -35,6 +35,7 @@
import org.apache.lucene.index.FilterDirectoryReader;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.store.AlreadyClosedException;
@@ -2326,8 +2327,7 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
// Case 3: MatchAllQuery and aggregations present, should use star tree if possible
searchService = getInstanceFromNode(SearchService.class);
sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.query(new MatchAllQueryBuilder());
- sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ sourceBuilder.size(0).query(new MatchAllQueryBuilder()).aggregation(AggregationBuilders.max("test").field("field"));
request.source(sourceBuilder);
try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
SearchContext context = searchService.createContext(reader, request, null, true);
@@ -2338,8 +2338,10 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
// Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
searchService = getInstanceFromNode(SearchService.class);
sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.query(new MatchAllQueryBuilder());
- sourceBuilder.aggregation(AggregationBuilders.max("test").field("field")).trackTotalHitsUpTo(1000);
+ sourceBuilder.size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"))
+ .trackTotalHitsUpTo(1000);
request.source(sourceBuilder);
try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
SearchContext context = searchService.createContext(reader, request, null, true);
@@ -2350,8 +2352,7 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
// Case 5: TermQuery and aggregations present, should use star tree
searchService = getInstanceFromNode(SearchService.class);
sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.query(new TermQueryBuilder("sndv", 1));
- sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ sourceBuilder.size(0).query(new TermQueryBuilder("sndv", 1)).aggregation(AggregationBuilders.max("test").field("field"));
request.source(sourceBuilder);
try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
SearchContext context = searchService.createContext(reader, request, null, true);
@@ -2362,7 +2363,7 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
// Case 6: No query, metric aggregations present, should use star tree
searchService = getInstanceFromNode(SearchService.class);
sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ sourceBuilder.size(0).aggregation(AggregationBuilders.max("test").field("field"));
request.source(sourceBuilder);
try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
SearchContext context = searchService.createContext(reader, request, null, true);
@@ -2370,6 +2371,18 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
}
searchService.doStop();
+ // Case 7: TermQuery and aggregations present, size !=0, should not use star tree
+ searchService = getInstanceFromNode(SearchService.class);
+ sourceBuilder = new SearchSourceBuilder();
+ sourceBuilder.query(new TermQueryBuilder("sndv", 1));
+ sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
+ request.source(sourceBuilder);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(IndexOrDocValuesQuery.class));
+ }
+ searchService.doStop();
+
searchService.doClose();
client().admin()
.cluster()
From 414d45a6d6cf84741ad0c5203f560398f174e6c6 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 4 Sep 2024 23:34:45 -0700
Subject: [PATCH 16/35] fix npe
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
server/src/main/java/org/opensearch/search/SearchService.java | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index a1cb51f098a2a..0626e2bdc2f90 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -1544,7 +1544,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch()));
}
- if (this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
+ if (this.indicesService.getCompositeIndexSettings() != null
+ && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
&& StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) {
try {
OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source);
From 99b5b67587f254a3f0b3715bfe2b0e0cc4db82dd Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Thu, 5 Sep 2024 00:23:20 -0700
Subject: [PATCH 17/35] search service test refactoring
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/search/SearchServiceTests.java | 117 ++++++------------
1 file changed, 41 insertions(+), 76 deletions(-)
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index 029e3fa302fec..8b8e038a41360 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -2271,11 +2271,7 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
- client().admin()
- .cluster()
- .prepareUpdateSettings()
- .setTransientSettings(Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), true).build())
- .execute();
+ setStarTreeIndexSetting("true");
Settings settings = Settings.builder()
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
@@ -2288,14 +2284,10 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
.setSettings(settings)
.setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
createIndex("test", builder);
+
IndicesService indicesService = getInstanceFromNode(IndicesService.class);
IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
IndexShard indexShard = indexService.getShard(0);
-
- SearchService searchService = getInstanceFromNode(SearchService.class);
-
- // Case 1: No query or aggregations, should not use star tree
- SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
ShardSearchRequest request = new ShardSearchRequest(
OriginalIndices.NONE,
new SearchRequest().allowPartialSearchResults(true),
@@ -2307,89 +2299,62 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
null,
null
);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, randomBoolean())) {
- SearchContext context = searchService.createContext(reader, request, null, false);
- assertFalse(context.query() instanceof StarTreeQuery);
- searchService.doStop();
- }
+
+ // Case 1: No query or aggregations, should not use star tree
+ SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+ assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
// Case 2: MatchAllQuery present but no aggregations, should not use star tree
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.query(new MatchAllQueryBuilder());
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, randomBoolean());
- assertThat(context.query(), instanceOf(MatchAllDocsQuery.class));
- searchService.doStop();
- }
+ sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
+ assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
- // Case 3: MatchAllQuery and aggregations present, should use star tree if possible
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.size(0).query(new MatchAllQueryBuilder()).aggregation(AggregationBuilders.max("test").field("field"));
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(StarTreeQuery.class));
- searchService.doStop();
- }
+ // Case 3: MatchAllQuery and aggregations present, should use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ assertQueryType(request, sourceBuilder, StarTreeQuery.class);
// Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.size(0)
+ sourceBuilder = new SearchSourceBuilder().size(0)
.query(new MatchAllQueryBuilder())
.aggregation(AggregationBuilders.max("test").field("field"))
.trackTotalHitsUpTo(1000);
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(MatchAllDocsQuery.class));
- }
- searchService.doStop();
+ assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
// Case 5: TermQuery and aggregations present, should use star tree
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.size(0).query(new TermQueryBuilder("sndv", 1)).aggregation(AggregationBuilders.max("test").field("field"));
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(StarTreeQuery.class));
- }
- searchService.doStop();
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ assertQueryType(request, sourceBuilder, StarTreeQuery.class);
// Case 6: No query, metric aggregations present, should use star tree
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.size(0).aggregation(AggregationBuilders.max("test").field("field"));
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(StarTreeQuery.class));
- }
- searchService.doStop();
+ sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
+ assertQueryType(request, sourceBuilder, StarTreeQuery.class);
- // Case 7: TermQuery and aggregations present, size !=0, should not use star tree
- searchService = getInstanceFromNode(SearchService.class);
- sourceBuilder = new SearchSourceBuilder();
- sourceBuilder.query(new TermQueryBuilder("sndv", 1));
- sourceBuilder.aggregation(AggregationBuilders.max("test").field("field"));
- request.source(sourceBuilder);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(IndexOrDocValuesQuery.class));
- }
- searchService.doStop();
+ // Case 7: TermQuery and aggregations present, size != 0, should not use star tree
+ sourceBuilder = new SearchSourceBuilder().query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ assertQueryType(request, sourceBuilder, IndexOrDocValuesQuery.class);
+
+ setStarTreeIndexSetting(null);
+ }
- searchService.doClose();
+ private void setStarTreeIndexSetting(String value) throws IOException {
client().admin()
.cluster()
.prepareUpdateSettings()
- .setTransientSettings(
- Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), (String) null).build()
- )
+ .setTransientSettings(Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), value).build())
.execute();
}
+
+ private void assertQueryType(ShardSearchRequest request, SearchSourceBuilder sourceBuilder, Class<?> expectedQueryClass)
+ throws IOException {
+ request.source(sourceBuilder);
+ SearchService searchService = getInstanceFromNode(SearchService.class);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ assertThat(context.query(), instanceOf(expectedQueryClass));
+ searchService.doStop();
+ }
+ }
}
From aef28f5b790c63d2c6359cf5956d36ee1cd996e6 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Thu, 5 Sep 2024 01:32:26 -0700
Subject: [PATCH 18/35] add changelog
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
CHANGELOG.md | 1 +
1 file changed, 1 insertion(+)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d4b6b848f5642..6f380e83589ab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
- Latency and Memory allocation improvements to Multi Term Aggregation queries ([#14993](https://github.com/opensearch-project/OpenSearch/pull/14993))
- Flat object field use IndexOrDocValuesQuery to optimize query ([#14383](https://github.com/opensearch-project/OpenSearch/issues/14383))
- Add method to return dynamic SecureTransportParameters from SecureTransportSettingsProvider interface ([#16387](https://github.com/opensearch-project/OpenSearch/pull/16387)
+- [Star Tree - Search] Add support for metric aggregations with/without term query ([15289](https://github.com/opensearch-project/OpenSearch/pull/15289))
### Dependencies
- Bump `com.azure:azure-identity` from 1.13.0 to 1.13.2 ([#15578](https://github.com/opensearch-project/OpenSearch/pull/15578))
From 03f15fb3ae5e043f392a96bd59c2b53e262e1f0e Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 18 Sep 2024 17:32:55 -0700
Subject: [PATCH 19/35] temp
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/common/util/FeatureFlags.java | 2 +-
.../startree/utils/StarTreeQueryHelper.java | 82 ++++++++++--
.../SortedNumericStarTreeValuesIterator.java | 4 +
.../iterator/StarTreeValuesIterator.java | 2 +-
.../search/DefaultSearchContext.java | 33 +++++
.../org/opensearch/search/SearchService.java | 9 +-
.../aggregations/metrics/AvgAggregator.java | 109 ++++++++--------
.../aggregations/metrics/MaxAggregator.java | 51 +++-----
.../aggregations/metrics/MinAggregator.java | 51 +++-----
.../aggregations/metrics/SumAggregator.java | 53 ++------
.../metrics/ValueCountAggregator.java | 36 ++----
.../search/internal/SearchContext.java | 18 ++-
.../startree/OriginalOrStarTreeQuery.java | 63 ---------
.../search/startree/StarTreeFilter.java | 42 +++---
.../search/startree/StarTreeQuery.java | 122 ------------------
.../search/startree/StarTreeQueryContext.java | 48 +++++++
16 files changed, 309 insertions(+), 416 deletions(-)
delete mode 100644 server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
delete mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
create mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index 6df68013a8119..e663d8429da13 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 351f38afff7fd..a1ec6cf7712c0 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -10,6 +10,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
+import org.apache.lucene.search.CollectionTerminatedException;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite.CompositeIndexReader;
@@ -17,22 +18,28 @@
import org.opensearch.index.compositeindex.datacube.Metric;
import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.aggregations.LeafBucketCollectorBase;
import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
+import org.opensearch.search.startree.StarTreeFilter;
+import org.opensearch.search.startree.StarTreeQueryContext;
import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.function.Consumer;
import java.util.stream.Collectors;
/**
@@ -43,6 +50,8 @@
*/
public class StarTreeQueryHelper {
+ private static Map<LeafReaderContext, StarTreeValues> starTreeValuesMap = new HashMap<>();
+
/**
* Checks if the search context can be supported by star-tree
*/
@@ -64,7 +73,12 @@ public static boolean isStarTreeSupported(SearchContext context, boolean trackTo
* Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
* Returns null if the query cannot be supported.
*/
- public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext context, SearchSourceBuilder source) throws IOException {
+
+ /**
+ * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
+ * Returns null if the query cannot be supported.
+ */
+ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
// Current implementation assumes only single star-tree is supported
CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
.getCompositeFieldTypes()
@@ -75,8 +89,12 @@ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext c
compositeMappedFieldType.getCompositeIndexType()
);
- StarTreeQuery starTreeQuery = StarTreeQueryHelper.toStarTreeQuery(starTree, compositeMappedFieldType, source.query());
- if (starTreeQuery == null) {
+ StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.toStarTreeQueryContext(
+ starTree,
+ compositeMappedFieldType,
+ source.query()
+ );
+ if (starTreeQueryContext == null) {
return null;
}
@@ -86,10 +104,10 @@ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext c
}
}
- return new OriginalOrStarTreeQuery(starTreeQuery, context.query());
+ return starTreeQueryContext;
}
- private static StarTreeQuery toStarTreeQuery(
+ private static StarTreeQueryContext toStarTreeQueryContext(
CompositeIndexFieldInfo starTree,
CompositeDataCubeFieldType compositeIndexFieldInfo,
QueryBuilder queryBuilder
@@ -110,7 +128,7 @@ private static StarTreeQuery toStarTreeQuery(
return null;
}
- return new StarTreeQuery(starTree, queryMap);
+ return new StarTreeQueryContext(starTree, queryMap);
}
/**
@@ -151,13 +169,11 @@ private static boolean validateStarTreeMetricSuport(
}
public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context) {
- if (context.query() instanceof StarTreeQuery) {
- return ((StarTreeQuery) context.query()).getStarTree();
- }
- return null;
+ StarTreeQueryContext starTreeQueryContext = context.getStarTreeQueryContext();
+ return (starTreeQueryContext != null) ? starTreeQueryContext.getStarTree() : null;
}
- public static StarTreeValues getStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException {
+ public static StarTreeValues computeStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException {
SegmentReader reader = Lucene.segmentReader(context.reader());
if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) {
return null;
@@ -165,4 +181,44 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos
CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
}
+
+ public static LeafBucketCollector getStarTreeLeafCollector(
+ SearchContext context,
+ ValuesSource.Numeric valuesSource,
+ LeafReaderContext ctx,
+ LeafBucketCollector sub,
+ CompositeIndexFieldInfo starTree,
+ String metric,
+ Consumer<Long> valueConsumer,
+ Runnable finalConsumer
+ ) throws IOException {
+ StarTreeValues starTreeValues = context.getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
+
+ assert starTreeValues != null;
+ SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+ metricName
+ );
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
+ StarTreeValuesIterator result = filter.getStarTreeResult();
+
+ int entryId;
+ while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ if (valuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ int count = valuesIterator.valuesCount();
+ for (int i = 0; i < count; i++) {
+ long value = valuesIterator.nextValue();
+ valueConsumer.accept(value); // Apply the operation (max, sum, etc.)
+ }
+ }
+ }
+ finalConsumer.run();
+ return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
+ @Override
+ public void collect(int doc, long bucket) {
+ throw new CollectionTerminatedException();
+ }
+ };
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
index 27afdf1479b4e..44f9545ce4f7f 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
@@ -29,4 +29,8 @@ public SortedNumericStarTreeValuesIterator(DocIdSetIterator docIdSetIterator) {
public long nextValue() throws IOException {
return ((SortedNumericDocValues) docIdSetIterator).nextValue();
}
+
+ public int valuesCount() throws IOException {
+ return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
index 32866f3e50092..454e5b393973f 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
@@ -21,7 +21,7 @@
* @opensearch.experimental
*/
@ExperimentalApi
-public abstract class StarTreeValuesIterator {
+public class StarTreeValuesIterator {
public static final int NO_MORE_ENTRIES = Integer.MAX_VALUE;
protected final DocIdSetIterator docIdSetIterator;
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
index 74a7482d975df..204e52c32342c 100644
--- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -34,6 +34,7 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
+import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.BoostQuery;
@@ -56,6 +57,8 @@
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
@@ -98,6 +101,7 @@
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.slice.SliceBuilder;
import org.opensearch.search.sort.SortAndFormats;
+import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
import java.io.IOException;
@@ -115,6 +119,7 @@
import java.util.function.Function;
import java.util.function.LongSupplier;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.computeStarTreeValues;
import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD;
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_MODE;
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
@@ -176,6 +181,7 @@ final class DefaultSearchContext extends SearchContext {
private SliceBuilder sliceBuilder;
private SearchShardTask task;
private final Version minNodeVersion;
+ private StarTreeQueryContext starTreeQueryContext;
/**
* The original query as sent by the user without the types and aliases
@@ -270,6 +276,7 @@ final class DefaultSearchContext extends SearchContext {
this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
+ this.starTreeValuesMap = new HashMap<>();
}
@Override
@@ -1147,4 +1154,30 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() {
}
return false;
}
+
+ @Override
+ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
+ this.starTreeQueryContext = starTreeQueryContext;
+ return this;
+ }
+
+ @Override
+ public StarTreeQueryContext getStarTreeQueryContext() {
+ return this.starTreeQueryContext;
+ }
+
+ @Override
+ public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+ if (this.starTreeValuesMap.containsKey(ctx)) {
+ logger.info("Used cached values");
+ return starTreeValuesMap.get(ctx);
+
+ } else {
+ logger.info("not using cache");
+ }
+
+ StarTreeValues starTreeValues = computeStarTreeValues(ctx, starTree);
+ starTreeValuesMap.put(ctx, starTreeValues);
+ return starTreeValues;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 0626e2bdc2f90..49f90a9f4c76c 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -84,7 +84,6 @@
import org.opensearch.index.query.InnerHitContextBuilder;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.MatchNoneQueryBuilder;
-import org.opensearch.index.query.ParsedQuery;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryRewriteContext;
import org.opensearch.index.query.QueryShardContext;
@@ -139,7 +138,7 @@
import org.opensearch.search.sort.SortAndFormats;
import org.opensearch.search.sort.SortBuilder;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.Suggest;
import org.opensearch.search.suggest.completion.CompletionSuggestion;
import org.opensearch.tasks.TaskResourceTrackingService;
@@ -1548,9 +1547,9 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
&& this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
&& StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) {
try {
- OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source);
- if (parsedQuery != null) {
- context.parsedQuery(new ParsedQuery(parsedQuery));
+ StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source);
+ if (starTreeQueryContext != null) {
+ context.starTreeQueryContext(starTreeQueryContext);
logger.debug("can use star tree");
} else {
logger.debug("cannot use star tree");
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index d028f38be1b61..017bb91871b0e 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -32,17 +32,12 @@
package org.opensearch.search.aggregations.metrics;
import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.common.util.LongArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -56,7 +51,6 @@
import java.io.IOException;
import java.util.Map;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -104,7 +98,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ // return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
@@ -144,56 +138,57 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
- throws IOException {
- final BigArrays bigArrays = context.bigArrays();
- final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
-
- StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.SUM.getTypeName()
- );
- assert starTreeValues != null;
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
-
- String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.VALUE_COUNT.getTypeName()
- );
- SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
-
- return new LeafBucketCollectorBase(sub, values) {
- @Override
- public void collect(int doc, long bucket) throws IOException {
- counts = bigArrays.grow(counts, bucket + 1);
- sums = bigArrays.grow(sums, bucket + 1);
- compensations = bigArrays.grow(compensations, bucket + 1);
-
- if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
- final long valueCount = values.docValueCount();
- counts.increment(bucket, countValues.nextValue());
- // Compute the sum of double values with Kahan summation algorithm which is more
- // accurate than naive summation.
- double sum = sums.get(bucket);
- double compensation = compensations.get(bucket);
-
- kahanSummation.reset(sum, compensation);
-
- for (int i = 0; i < valueCount; i++) {
- double value = NumericUtils.sortableLongToDouble(values.nextValue());
- kahanSummation.add(value);
- }
-
- sums.set(bucket, kahanSummation.value());
- compensations.set(bucket, kahanSummation.delta());
- }
- }
- };
- }
+ // private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo
+ // starTree)
+ // throws IOException {
+ // final BigArrays bigArrays = context.bigArrays();
+ // final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ //
+ // StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ // String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ // String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ // starTree.getField(),
+ // fieldName,
+ // MetricStat.SUM.getTypeName()
+ // );
+ // assert starTreeValues != null;
+ // SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
+ //
+ // String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ // starTree.getField(),
+ // fieldName,
+ // MetricStat.VALUE_COUNT.getTypeName()
+ // );
+ // SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
+ //
+ // return new LeafBucketCollectorBase(sub, values) {
+ // @Override
+ // public void collect(int doc, long bucket) throws IOException {
+ // counts = bigArrays.grow(counts, bucket + 1);
+ // sums = bigArrays.grow(sums, bucket + 1);
+ // compensations = bigArrays.grow(compensations, bucket + 1);
+ //
+ // if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
+ // final long valueCount = values.docValueCount();
+ // counts.increment(bucket, countValues.nextValue());
+ // // Compute the sum of double values with Kahan summation algorithm which is more
+ // // accurate than naive summation.
+ // double sum = sums.get(bucket);
+ // double compensation = compensations.get(bucket);
+ //
+ // kahanSummation.reset(sum, compensation);
+ //
+ // for (int i = 0; i < valueCount; i++) {
+ // double value = NumericUtils.sortableLongToDouble(values.nextValue());
+ // kahanSummation.add(value);
+ // }
+ //
+ // sums.set(bucket, kahanSummation.value());
+ // compensations.set(bucket, kahanSummation.delta());
+ // }
+ // }
+ // };
+ // }
@Override
public double metric(long owningBucketOrd) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 5f0c5e242a453..c7e697ad45d8f 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -34,7 +34,6 @@
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PointValues;
-import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
@@ -44,8 +43,7 @@
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.NumericDoubleValues;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
@@ -61,9 +59,9 @@
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -132,7 +130,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ System.out.println("max star tree");
+ return getStarTreeCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
@@ -162,37 +161,21 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.MAX.getTypeName()
+ AtomicReference<Double> max = new AtomicReference<>(Double.NEGATIVE_INFINITY);
+ return StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.SUM.getTypeName(),
+ value -> {
+ max.set(Math.max(max.get(), (NumericUtils.sortableLongToDouble(value))));
+ },
+ () -> maxes.set(0, max.get())
);
- assert starTreeValues != null;
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
-
- final BigArrays bigArrays = context.bigArrays();
- final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
- return new LeafBucketCollectorBase(sub, allValues) {
-
- @Override
- public void collect(int doc, long bucket) throws IOException {
- if (bucket >= maxes.size()) {
- long from = maxes.size();
- maxes = bigArrays.grow(maxes, bucket + 1);
- maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY);
- }
- if (values.advanceExact(doc)) {
- final double value = NumericUtils.sortableLongToDouble(values.nextValue());
- double max = maxes.get(bucket);
- max = Math.max(max, value);
- maxes.set(bucket, max);
- }
- }
- };
}
@Override
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 81246a192555c..188c615a26ff2 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -34,7 +34,6 @@
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PointValues;
-import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.Bits;
@@ -44,8 +43,7 @@
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.NumericDoubleValues;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
@@ -60,9 +58,9 @@
import java.io.IOException;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -131,7 +129,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ System.out.println("min star tree");
+ return getStarTreeCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
@@ -159,37 +158,21 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.MIN.getTypeName()
+ AtomicReference<Double> min = new AtomicReference<>(Double.POSITIVE_INFINITY);
+ return StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.SUM.getTypeName(),
+ value -> {
+ min.set(Math.min(min.get(), (NumericUtils.sortableLongToDouble(value))));
+ },
+ () -> mins.set(0, min.get())
);
- assert starTreeValues != null;
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
-
- final BigArrays bigArrays = context.bigArrays();
- final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx);
- return new LeafBucketCollectorBase(sub, allValues) {
-
- @Override
- public void collect(int doc, long bucket) throws IOException {
- if (bucket >= mins.size()) {
- long from = mins.size();
- mins = bigArrays.grow(mins, bucket + 1);
- mins.fill(from, mins.size(), Double.POSITIVE_INFINITY);
- }
- if (values.advanceExact(doc)) {
- final double value = NumericUtils.sortableLongToDouble(values.nextValue());
- double min = mins.get(bucket);
- min = Math.min(min, value);
- mins.set(bucket, min);
- }
- }
- };
}
@Override
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index c205c9abb7bee..617ee6a939e6c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -32,7 +32,6 @@
package org.opensearch.search.aggregations.metrics;
import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
@@ -40,8 +39,7 @@
import org.opensearch.common.util.DoubleArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -55,7 +53,6 @@
import java.io.IOException;
import java.util.Map;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -101,7 +98,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ return getStarTreeCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
@@ -136,43 +133,19 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.SUM.getTypeName()
- );
- assert starTreeValues != null;
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
-
- final BigArrays bigArrays = context.bigArrays();
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
-
- return new LeafBucketCollectorBase(sub, values) {
- @Override
- public void collect(int doc, long bucket) throws IOException {
- sums = bigArrays.grow(sums, bucket + 1);
- compensations = bigArrays.grow(compensations, bucket + 1);
-
- if (values.advanceExact(doc)) {
- final int valuesCount = values.docValueCount();
- double sum = sums.get(bucket);
- double compensation = compensations.get(bucket);
- kahanSummation.reset(sum, compensation);
-
- for (int i = 0; i < valuesCount; i++) {
- double value = NumericUtils.sortableLongToDouble(values.nextValue());
- kahanSummation.add(value);
- }
-
- compensations.set(bucket, kahanSummation.delta());
- sums.set(bucket, kahanSummation.value());
- }
- }
- };
+ return StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.SUM.getTypeName(),
+ value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),
+ () -> sums.set(0, kahanSummation.value())
+ );
}
@Override
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
index 0d9a145b2b8ad..a156ec49983fa 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java
@@ -39,8 +39,7 @@
import org.opensearch.common.util.LongArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.MetricStat;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.MultiGeoPointValues;
import org.opensearch.index.fielddata.SortedBinaryDocValues;
import org.opensearch.search.aggregations.Aggregator;
@@ -54,7 +53,6 @@
import java.io.IOException;
import java.util.Map;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -98,7 +96,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ return getStarTreeCollector(ctx, sub, supportedStarTree);
}
final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx);
@@ -140,28 +138,18 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- starTree.getField(),
- fieldName,
- MetricStat.VALUE_COUNT.getTypeName()
+ return StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ (ValuesSource.Numeric) valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.VALUE_COUNT.getTypeName(),
+ value -> counts.increment(0, value),
+ () -> {}
);
- assert starTreeValues != null;
- SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName);
- final BigArrays bigArrays = context.bigArrays();
-
- return new LeafBucketCollectorBase(sub, values) {
- @Override
- public void collect(int doc, long bucket) throws IOException {
- counts = bigArrays.grow(counts, bucket + 1);
- if (values.advanceExact(doc)) {
- counts.increment(bucket, values.nextValue());
- }
- }
- };
}
@Override
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index 5357206e8c117..ed3cbde61ce74 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -31,6 +31,7 @@
package org.opensearch.search.internal;
+import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.FieldDoc;
@@ -44,6 +45,8 @@
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.util.BigArrays;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.ObjectMapper;
@@ -76,8 +79,10 @@
import org.opensearch.search.query.ReduceableSearchResult;
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.sort.SortAndFormats;
+import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
+import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
@@ -124,7 +129,7 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private final List<Releasable> releasables = new CopyOnWriteArrayList<>();
private final AtomicBoolean closed = new AtomicBoolean(false);
private InnerHitsContext innerHitsContext;
-
+ protected volatile Map<LeafReaderContext, StarTreeValues> starTreeValuesMap;
private volatile boolean searchTimedOut;
protected SearchContext() {}
@@ -531,4 +536,15 @@ public boolean keywordIndexOrDocValuesEnabled() {
return false;
}
+ public StarTreeQueryContext getStarTreeQueryContext() {
+ return null;
+ }
+
+ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
+ return this;
+ }
+
+ public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+ return null;
+ }
}
diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
deleted file mode 100644
index 4806265888a2c..0000000000000
--- a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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;
-
-import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.Query;
-import org.apache.lucene.search.QueryVisitor;
-
-import java.io.IOException;
-import java.util.Objects;
-
-/**
- * Preserves star-tree queries which can be used along with original query
- * Decides which star-tree query to use (or not) based on cost factors
- *
- * @opensearch.experimental
- */
-public class OriginalOrStarTreeQuery extends Query {
-
- private final StarTreeQuery starTreeQuery;
- private final Query originalQuery;
-
- public OriginalOrStarTreeQuery(StarTreeQuery starTreeQuery, Query originalQuery) {
- this.starTreeQuery = starTreeQuery;
- this.originalQuery = originalQuery;
- }
-
- @Override
- public String toString(String s) {
- return originalQuery.toString(s);
- }
-
- @Override
- public void visit(QueryVisitor queryVisitor) {}
-
- @Override
- public boolean equals(Object o) {
- return sameClassAs(o) && equalsTo(getClass().cast(o));
- }
-
- private boolean equalsTo(OriginalOrStarTreeQuery other) {
- return starTreeQuery.equals(other.starTreeQuery) && originalQuery.equals(other.originalQuery);
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(classHash(), starTreeQuery, originalQuery, starTreeQuery);
- }
-
- @Override
- public Query rewrite(IndexSearcher indexSearcher) throws IOException {
- if (indexSearcher.getIndexReader().hasDeletions()) {
- return originalQuery;
- }
- return starTreeQuery;
- }
-}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 7b48ecd4561c7..493bdf3ca143e 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -10,13 +10,13 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
-import org.apache.lucene.index.SortedNumericDocValues;
-import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.DocIdSetBuilder;
import org.opensearch.index.compositeindex.datacube.Dimension;
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.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import java.io.IOException;
import java.util.ArrayDeque;
@@ -38,7 +38,7 @@
* @opensearch.experimental
* @opensearch.internal
*/
-class StarTreeFilter {
+public class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
private final Map<String, Long> queryMap;
@@ -53,49 +53,49 @@ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> pr
/**
* <ul>
* <li>First go over the star tree and try to match as many dimensions as possible
- * <li>For the remaining columns, use doc values indexes to match them
+ * <li>For the remaining columns, use star-tree doc values to match them
* </ul>
*/
- public DocIdSetIterator getStarTreeResult() throws IOException {
+ public StarTreeValuesIterator getStarTreeResult() throws IOException {
StarTreeResult starTreeResult = traverseStarTree();
- List<DocIdSetIterator> andIterators = new ArrayList<>();
- andIterators.add(starTreeResult._matchedDocIds.build().iterator());
- DocIdSetIterator docIdSetIterator = andIterators.get(0);
+ List<StarTreeValuesIterator> andIterators = new ArrayList<>();
+ andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator()));
+ StarTreeValuesIterator starTreeValuesIterator = andIterators.get(0);
// No matches, return
if (starTreeResult.maxMatchedDoc == -1) {
- return docIdSetIterator;
+ return starTreeValuesIterator;
}
for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
- SortedNumericDocValues ndv = (SortedNumericDocValues) this.starTreeValues.getDimensionDocIdSetIterator(
+ SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) this.starTreeValues.getDimensionValuesIterator(
remainingPredicateColumn
);
- List<Integer> docIds = new ArrayList<>();
+ List<Integer> entryIds = new ArrayList<>();
long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly
- while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) {
- int docID = docIdSetIterator.docID();
- if (ndv.advanceExact(docID)) {
- final int valuesCount = ndv.docValueCount();
+ while (starTreeValuesIterator.nextEntry() != NO_MORE_DOCS) {
+ int entryId = starTreeValuesIterator.entryId();
+ if (ndv.advance(entryId) > 0) {
+ final int valuesCount = ndv.valuesCount();
for (int i = 0; i < valuesCount; i++) {
long value = ndv.nextValue();
// Directly compare value with queryValue
if (value == queryValue) {
- docIds.add(docID);
+ entryIds.add(entryId);
break;
}
}
}
}
- DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size());
- for (int docID : docIds) {
- adder.add(docID);
+ DocIdSetBuilder.BulkAdder adder = builder.grow(entryIds.size());
+ for (int entryId : entryIds) {
+ adder.add(entryId);
}
- docIdSetIterator = builder.build().iterator();
+ starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator());
}
- return docIdSetIterator;
+ return starTreeValuesIterator;
}
/**
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
deleted file mode 100644
index e8354b0415ac1..0000000000000
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
- * 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;
-
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SegmentReader;
-import org.apache.lucene.search.ConstantScoreScorer;
-import org.apache.lucene.search.ConstantScoreWeight;
-import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.Query;
-import org.apache.lucene.search.QueryVisitor;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
-import org.apache.lucene.search.Weight;
-import org.opensearch.common.lucene.Lucene;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.codec.composite.CompositeIndexReader;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-
-import java.io.IOException;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-
-/**
- * Query class for querying star tree data structure.
- *
- * @opensearch.experimental
- */
-public class StarTreeQuery extends Query {
-
- /**
- * Star tree field info
- * This is used to get the star tree data structure
- */
- private final CompositeIndexFieldInfo starTree;
-
- /**
- * Map of field name to a value to be queried for that field
- * This is used to filter the data based on the query
- */
- private final Map<String, Long> queryMap;
-
- public StarTreeQuery(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
- this.starTree = starTree;
- this.queryMap = queryMap;
- }
-
- @Override
- public void visit(QueryVisitor visitor) {}
-
- @Override
- public boolean equals(Object obj) {
- return sameClassAs(obj) && equalsTo(getClass().cast(obj));
- }
-
- private boolean equalsTo(StarTreeQuery other) {
- return starTree.equals(other.starTree) && queryMap != null && queryMap.equals(other.queryMap);
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(classHash(), starTree, queryMap);
- }
-
- @Override
- public String toString(String field) {
- StringBuilder sb = new StringBuilder();
- sb.append(getClass().getSimpleName());
- sb.append("(");
- sb.append(this.starTree);
- if (queryMap != null) {
- sb.append(", ");
- sb.append(queryMap);
- sb.append(")");
- }
- return sb.toString();
- }
-
- @Override
- public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
- return new ConstantScoreWeight(this, boost) {
- @Override
- public Scorer scorer(LeafReaderContext context) throws IOException {
- StarTreeValues starTreeValues = getStarTreeValues(context);
- if (starTreeValues == null) {
- return null;
- }
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, queryMap);
- DocIdSetIterator result = filter.getStarTreeResult();
- return new ConstantScoreScorer(this, score(), scoreMode, result);
- }
-
- @Override
- public boolean isCacheable(LeafReaderContext ctx) {
- return false;
- }
-
- private StarTreeValues getStarTreeValues(LeafReaderContext ctx) throws IOException {
- SegmentReader reader = Lucene.segmentReader(ctx.reader());
- CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
- List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
- if (compositeIndexFields != null && !compositeIndexFields.isEmpty()) {
- return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
- } else {
- return null;
- }
- }
- };
- }
-
- public CompositeIndexFieldInfo getStarTree() {
- return starTree;
- }
-}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
new file mode 100644
index 0000000000000..2b12c9bba604f
--- /dev/null
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+
+import java.util.Map;
+
+/**
+ * Query class for querying star tree data structure.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class StarTreeQueryContext {
+
+ /**
+ * Star tree field info
+ * This is used to get the star tree data structure
+ */
+ private final CompositeIndexFieldInfo starTree;
+
+ /**
+ * Map of field name to a value to be queried for that field
+ * This is used to filter the data based on the query
+ */
+ private final Map<String, Long> queryMap;
+
+ public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
+ this.starTree = starTree;
+ this.queryMap = queryMap;
+ }
+
+ public CompositeIndexFieldInfo getStarTree() {
+ return starTree;
+ }
+
+ public Map<String, Long> getQueryMap() {
+ return queryMap;
+ }
+}
From 0979f5b0fa917c4b550099c04ceee53ed100d355 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Thu, 19 Sep 2024 03:29:34 -0700
Subject: [PATCH 20/35] temp temp
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 7 +-
.../search/DefaultSearchContext.java | 19 +-
.../aggregations/metrics/AvgAggregator.java | 201 +++++++++++++-----
.../aggregations/metrics/MaxAggregator.java | 1 -
.../aggregations/metrics/MinAggregator.java | 1 -
.../search/internal/SearchContext.java | 5 +-
.../search/startree/StarTreeFilter.java | 3 +
7 files changed, 169 insertions(+), 68 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index a1ec6cf7712c0..04a98e82a97fd 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -173,7 +173,7 @@ public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context
return (starTreeQueryContext != null) ? starTreeQueryContext.getStarTree() : null;
}
- public static StarTreeValues computeStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException {
+ public static StarTreeValues getStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException {
SegmentReader reader = Lucene.segmentReader(context.reader());
if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) {
return null;
@@ -192,7 +192,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
Consumer<Long> valueConsumer,
Runnable finalConsumer
) throws IOException {
- StarTreeValues starTreeValues = context.getStarTreeValues(ctx, starTree);
+ StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
@@ -200,8 +200,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
metricName
);
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
- StarTreeValuesIterator result = filter.getStarTreeResult();
+ StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
int entryId;
while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
index 204e52c32342c..fd53fa4153b1e 100644
--- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -43,6 +43,7 @@
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
+import org.apache.lucene.util.BitSet;
import org.opensearch.Version;
import org.opensearch.action.search.SearchShardTask;
import org.opensearch.action.search.SearchType;
@@ -59,6 +60,8 @@
import org.opensearch.index.cache.bitset.BitsetFilterCache;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
@@ -101,6 +104,7 @@
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.slice.SliceBuilder;
import org.opensearch.search.sort.SortAndFormats;
+import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
@@ -119,7 +123,6 @@
import java.util.function.Function;
import java.util.function.LongSupplier;
-import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.computeStarTreeValues;
import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD;
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_MODE;
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
@@ -1167,17 +1170,15 @@ public StarTreeQueryContext getStarTreeQueryContext() {
}
@Override
- public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+ public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
if (this.starTreeValuesMap.containsKey(ctx)) {
- logger.info("Used cached values");
- return starTreeValuesMap.get(ctx);
- } else {
- logger.info("not using cache");
+ return starTreeValuesMap.get(ctx);
}
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, this.getStarTreeQueryContext().getQueryMap());
+ StarTreeValuesIterator result = filter.getStarTreeResult();
- StarTreeValues starTreeValues = computeStarTreeValues(ctx, starTree);
- starTreeValuesMap.put(ctx, starTreeValues);
- return starTreeValues;
+ starTreeValuesMap.put(ctx, result);
+ return result;
}
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 017bb91871b0e..f1b32fe8c9602 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -32,12 +32,20 @@
package org.opensearch.search.aggregations.metrics;
import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.DoubleArray;
import org.opensearch.common.util.LongArray;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -47,9 +55,11 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.StarTreeFilter;
import java.io.IOException;
import java.util.Map;
+import java.util.function.Consumer;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
@@ -138,57 +148,146 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- // private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo
- // starTree)
- // throws IOException {
- // final BigArrays bigArrays = context.bigArrays();
- // final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
- //
- // StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
- // String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
- // String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- // starTree.getField(),
- // fieldName,
- // MetricStat.SUM.getTypeName()
- // );
- // assert starTreeValues != null;
- // SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
- //
- // String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
- // starTree.getField(),
- // fieldName,
- // MetricStat.VALUE_COUNT.getTypeName()
- // );
- // SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
- //
- // return new LeafBucketCollectorBase(sub, values) {
- // @Override
- // public void collect(int doc, long bucket) throws IOException {
- // counts = bigArrays.grow(counts, bucket + 1);
- // sums = bigArrays.grow(sums, bucket + 1);
- // compensations = bigArrays.grow(compensations, bucket + 1);
- //
- // if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
- // final long valueCount = values.docValueCount();
- // counts.increment(bucket, countValues.nextValue());
- // // Compute the sum of double values with Kahan summation algorithm which is more
- // // accurate than naive summation.
- // double sum = sums.get(bucket);
- // double compensation = compensations.get(bucket);
- //
- // kahanSummation.reset(sum, compensation);
- //
- // for (int i = 0; i < valueCount; i++) {
- // double value = NumericUtils.sortableLongToDouble(values.nextValue());
- // kahanSummation.add(value);
- // }
- //
- // sums.set(bucket, kahanSummation.value());
- // compensations.set(bucket, kahanSummation.delta());
- // }
- // }
- // };
- // }
+ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ return StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.SUM.getTypeName(),
+ value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),
+ () -> sums.set(0, kahanSummation.value())
+ );
+ }
+
+// private LeafBucketCollector getStarTreeLeafCollector1(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+// throws IOException {
+// final BigArrays bigArrays = context.bigArrays();
+// final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+//
+// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+// String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+// starTree.getField(),
+// fieldName,
+// MetricStat.SUM.getTypeName()
+// );
+// assert starTreeValues != null;
+// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
+//
+// String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+// starTree.getField(),
+// fieldName,
+// MetricStat.VALUE_COUNT.getTypeName()
+// );
+// SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
+//
+// return new LeafBucketCollectorBase(sub, values) {
+// @Override
+// public void collect(int doc, long bucket) throws IOException {
+// counts = bigArrays.grow(counts, bucket + 1);
+// sums = bigArrays.grow(sums, bucket + 1);
+// compensations = bigArrays.grow(compensations, bucket + 1);
+//
+// if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
+// final long valueCount = values.docValueCount();
+// counts.increment(bucket, countValues.nextValue());
+// // Compute the sum of double values with Kahan summation algorithm which is more
+// // accurate than naive summation.
+// double sum = sums.get(bucket);
+// double compensation = compensations.get(bucket);
+//
+// kahanSummation.reset(sum, compensation);
+//
+// for (int i = 0; i < valueCount; i++) {
+// double value = NumericUtils.sortableLongToDouble(values.nextValue());
+// kahanSummation.add(value);
+// }
+//
+// sums.set(bucket, kahanSummation.value());
+// compensations.set(bucket, kahanSummation.delta());
+// }
+// }
+// };
+// }
+
+
+ public LeafBucketCollector getStarTreeCollector2(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ throws IOException {
+ final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.SUM.getTypeName(),
+ value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),
+ () -> sums.set(0, kahanSummation.value())
+ );
+ StarTreeQueryHelper.getStarTreeLeafCollector(
+ context,
+ valuesSource,
+ ctx,
+ sub,
+ starTree,
+ MetricStat.VALUE_COUNT.getTypeName(),
+ value -> counts.increment(0, value),
+ () -> {}
+ );
+ return LeafBucketCollector.NO_OP_COLLECTOR;
+ }
+
+ public LeafBucketCollector getStarTreeLeafCollector(
+ LeafReaderContext ctx,
+ LeafBucketCollector sub,
+ CompositeIndexFieldInfo starTree
+ ) throws IOException {
+ StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
+ String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+ String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.SUM.getTypeName()
+ );
+ String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ fieldName,
+ MetricStat.VALUE_COUNT.getTypeName()
+ );
+
+ assert starTreeValues != null;
+
+ final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+ sumMetricName
+ );
+ SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+ countMetricName
+ );
+ StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
+
+ int entryId;
+ while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ if (sumValuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ int count = sumValuesIterator.valuesCount();
+ for (int i = 0; i < count; i++) {
+ kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue()));
+ counts.increment(0, countValueIterator.nextValue());
+ }
+ }
+ }
+ sums.set(0, kahanSummation.value());
+ return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
+ @Override
+ public void collect(int doc, long bucket) {
+ throw new CollectionTerminatedException();
+ }
+ };
+ }
@Override
public double metric(long owningBucketOrd) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index c7e697ad45d8f..e68dcc4acb614 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -130,7 +130,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- System.out.println("max star tree");
return getStarTreeCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 188c615a26ff2..11f38020aa8a1 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -129,7 +129,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- System.out.println("min star tree");
return getStarTreeCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index ed3cbde61ce74..66b2ba9c98431 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -47,6 +47,7 @@
import org.opensearch.index.cache.bitset.BitsetFilterCache;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.ObjectMapper;
@@ -129,7 +130,7 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private final List<Releasable> releasables = new CopyOnWriteArrayList<>();
private final AtomicBoolean closed = new AtomicBoolean(false);
private InnerHitsContext innerHitsContext;
- protected volatile Map<LeafReaderContext, StarTreeValues> starTreeValuesMap;
+ protected volatile Map<LeafReaderContext, StarTreeValuesIterator> starTreeValuesMap;
private volatile boolean searchTimedOut;
protected SearchContext() {}
@@ -544,7 +545,7 @@ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryCont
return this;
}
- public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+ public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
return null;
}
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 493bdf3ca143e..fee32ba48e1f5 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -11,6 +11,7 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.util.DocIdSetBuilder;
+import org.apache.lucene.util.FixedBitSet;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
@@ -62,6 +63,7 @@ public StarTreeValuesIterator getStarTreeResult() throws IOException {
andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator()));
StarTreeValuesIterator starTreeValuesIterator = andIterators.get(0);
+ int length = 0;
// No matches, return
if (starTreeResult.maxMatchedDoc == -1) {
return starTreeValuesIterator;
@@ -93,6 +95,7 @@ public StarTreeValuesIterator getStarTreeResult() throws IOException {
for (int entryId : entryIds) {
adder.add(entryId);
}
+ length = entryIds.size();
starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator());
}
return starTreeValuesIterator;
From 94aabe5e99c5934acd14de3a3d683037f0a334a8 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 25 Sep 2024 00:50:23 -0700
Subject: [PATCH 21/35] temp temp temp
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 178 ++++++++++++++++--
.../search/DefaultSearchContext.java | 6 +-
.../aggregations/metrics/AvgAggregator.java | 22 ++-
.../aggregations/metrics/MaxAggregator.java | 2 +-
.../aggregations/metrics/MinAggregator.java | 2 +-
.../aggregations/metrics/SumAggregator.java | 3 +-
.../search/internal/SearchContext.java | 5 +-
.../search/startree/StarTreeFilter.java | 56 +++---
.../search/startree/StarTreeQueryContext.java | 31 +++
9 files changed, 253 insertions(+), 52 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 04a98e82a97fd..d516a886b6170 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -11,6 +11,8 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.CollectionTerminatedException;
+import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite.CompositeIndexReader;
@@ -28,14 +30,15 @@
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.LeafBucketCollector;
import org.opensearch.search.aggregations.LeafBucketCollectorBase;
+import org.opensearch.search.aggregations.metrics.CompensatedSum;
import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -50,8 +53,6 @@
*/
public class StarTreeQueryHelper {
- private static Map<LeafReaderContext, StarTreeValues> starTreeValuesMap = new HashMap<>();
-
/**
* Checks if the search context can be supported by star-tree
*/
@@ -69,10 +70,6 @@ public static boolean isStarTreeSupported(SearchContext context, boolean trackTo
return canUseStarTree;
}
- /**
- * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
- * Returns null if the query cannot be supported.
- */
/**
* Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
@@ -98,10 +95,14 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
return null;
}
+ boolean needCaching = context.aggregations().factories().getFactories().length > 1;
+// List<MetricInfo> metricInfos = new ArrayList<>();
for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
- if (validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
+ MetricStat metricStat = validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory);
+ if (metricStat == null) {
return null;
}
+// metricInfos.add(new )
}
return starTreeQueryContext;
@@ -150,10 +151,11 @@ private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder
return predicateMap;
}
- private static boolean validateStarTreeMetricSuport(
+ private static MetricStat validateStarTreeMetricSuport(
CompositeDataCubeFieldType compositeIndexFieldInfo,
AggregatorFactory aggregatorFactory
) {
+// List<MetricStat> metricStats = new ArrayList<>();
if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) {
String field;
Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
@@ -162,10 +164,12 @@ private static boolean validateStarTreeMetricSuport(
MetricStat metricStat = ((MetricAggregatorFactory) aggregatorFactory).getMetricStat();
field = ((MetricAggregatorFactory) aggregatorFactory).getField();
- return supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat);
- } else {
- return false;
+
+ if (supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(metricStat)) {
+ return metricStat;
+ }
}
+ return null;
}
public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context) {
@@ -200,19 +204,32 @@ public static LeafBucketCollector getStarTreeLeafCollector(
SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
metricName
);
- StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
+ // Obtain a FixedBitSet of matched document IDs
+ FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues); // Assuming this method gives a FixedBitSet
+
+ // Safety check: make sure the FixedBitSet is non-null and valid
+ if (matchedDocIds == null) {
+ throw new IllegalStateException("FixedBitSet is null");
+ }
+
+ int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
- int entryId;
- while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- if (valuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ // Iterate over the FixedBitSet
+ for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ // Advance to the bit (entryId) in the valuesIterator
+ if (valuesIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
int count = valuesIterator.valuesCount();
for (int i = 0; i < count; i++) {
long value = valuesIterator.nextValue();
- valueConsumer.accept(value); // Apply the operation (max, sum, etc.)
+ valueConsumer.accept(value); // Apply the consumer operation (e.g., max, sum)
}
}
}
+
+ // Call the final consumer after processing all entries
finalConsumer.run();
+
+ // Return a LeafBucketCollector that terminates collection
return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
@Override
public void collect(int doc, long bucket) {
@@ -220,4 +237,131 @@ public void collect(int doc, long bucket) {
}
};
}
+
+// public static LeafBucketCollector getStarTreeLeafCollectorNew(
+// SearchContext context,
+// ValuesSource.Numeric valuesSource,
+// LeafReaderContext ctx,
+// LeafBucketCollector sub,
+// CompositeIndexFieldInfo starTree,
+// String metric,
+// Consumer<Long> valueConsumer,
+// Runnable finalConsumer
+// ) throws IOException {
+// // Check in contextCache if the star-tree values are already computed
+// Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> cache = context.getStarTreeQueryContext().getLeafResultsCache();
+// if(cache != null) {
+//
+// if (cache.containsKey(ctx)) {
+// MetricInfo metricInfoMap = cache.get(ctx).get(metric);
+// finalConsumer.run();
+// }
+// }
+// else if (!cache.containsKey(ctx)) {
+// // TODO: fetch from cache
+//
+// } else {
+// // TODO: compute cache first
+// }
+//
+// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
+// String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
+//
+// assert starTreeValues != null;
+// List<SortedNumericStarTreeValuesIterator> valuesIterators;
+// SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+// metricName
+// );
+// StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
+//
+// int entryId;
+// while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+// for
+// if (valuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+// int count = valuesIterator.valuesCount();
+// for (int i = 0; i < count; i++) {
+// long value = valuesIterator.nextValue();
+// valueConsumer.accept(value); // Apply the operation (max, sum, etc.)
+// }
+// }
+// }
+// finalConsumer.run();
+// return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
+// @Override
+// public void collect(int doc, long bucket) {
+// throw new CollectionTerminatedException();
+// }
+// };
+// }
+//
+// public abstract class MetricInfo {
+// String metric;
+// MetricStat metricStat;
+//
+//
+//
+// MetricInfo (String metric, MetricStat metricStat) {
+// if (metricStat == MetricStat.SUM) {
+// return new SumMetricInfo(metric);
+// }
+// return null;
+// }
+//
+//
+// public abstract void valueConsumer(long value);
+//
+// public abstract <T extends Number> T getMetricValue();
+// }
+//
+// public class SumMetricInfo extends MetricInfo {
+// CompensatedSum compensatedSum;
+//
+// public SumMetricInfo(String metric) {
+// super(metric, MetricStat.SUM);
+// compensatedSum = new CompensatedSum(0,0);
+// }
+//
+// public void valueConsumer(long value) {
+// compensatedSum.add(NumericUtils.sortableLongToDouble(value));
+// }
+//
+// public Double getMetricValue() {
+// return compensatedSum.value();
+// }
+// }
+//
+// public static void computeLeafResultsCache(SearchContext context,
+// LeafReaderContext ctx,
+// CompositeIndexFieldInfo starTree,
+// List<MetricInfo> metricInfos) throws IOException {
+// Map<String, MetricInfo> leafCache = new HashMap<>();
+// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+// assert starTreeValues != null;
+// StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
+//
+// List<Integer> entryIdCache = new ArrayList<>();
+// int entryId;
+// while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+// entryIdCache.add(entryId);
+// }
+//
+// for (MetricInfo metricInfo : metricInfos) {
+// SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+// metricInfo.metric
+// );
+//
+// for (int cachedEntryId : entryIdCache) {
+// if (valuesIterator.advance(cachedEntryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+// int count = valuesIterator.valuesCount();
+// for (int i = 0; i < count; i++) {
+// long value = valuesIterator.nextValue();
+// metricInfo.valueConsumer(value);
+// }
+// }
+// }
+// leafCache.put(metricInfo.metric, metricInfo);
+// }
+// context.getStarTreeQueryContext().getLeafResultsCache().put(ctx, leafCache);
+// }
}
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
index fd53fa4153b1e..cfe2ad7681a58 100644
--- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -44,6 +44,7 @@
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.FixedBitSet;
import org.opensearch.Version;
import org.opensearch.action.search.SearchShardTask;
import org.opensearch.action.search.SearchType;
@@ -1170,13 +1171,12 @@ public StarTreeQueryContext getStarTreeQueryContext() {
}
@Override
- public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
+ public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
if (this.starTreeValuesMap.containsKey(ctx)) {
-
return starTreeValuesMap.get(ctx);
}
StarTreeFilter filter = new StarTreeFilter(starTreeValues, this.getStarTreeQueryContext().getQueryMap());
- StarTreeValuesIterator result = filter.getStarTreeResult();
+ FixedBitSet result = filter.getStarTreeResult();
starTreeValuesMap.put(ctx, result);
return result;
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index f1b32fe8c9602..ffd14976b57b0 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -34,6 +34,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
@@ -268,18 +269,29 @@ public LeafBucketCollector getStarTreeLeafCollector(
SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
countMetricName
);
- StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
+ FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues);
- int entryId;
- while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- if (sumValuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+
+ // Safety check: make sure the FixedBitSet is non-null and valid
+ if (matchedDocIds == null) {
+ throw new IllegalStateException("FixedBitSet is null");
+ }
+
+ int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
+
+ // Iterate over the FixedBitSet
+ for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ // Advance to the bit (entryId) in the valuesIterator
+ if (sumValuesIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES &&
+ countValueIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
int count = sumValuesIterator.valuesCount();
for (int i = 0; i < count; i++) {
kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue()));
- counts.increment(0, countValueIterator.nextValue());
+ counts.increment(0, countValueIterator.nextValue()); // Apply the consumer operation (e.g., max, sum)
}
}
}
+
sums.set(0, kahanSummation.value());
return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
@Override
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index e68dcc4acb614..338f12ea238e1 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -169,7 +169,7 @@ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucke
ctx,
sub,
starTree,
- MetricStat.SUM.getTypeName(),
+ MetricStat.MAX.getTypeName(),
value -> {
max.set(Math.max(max.get(), (NumericUtils.sortableLongToDouble(value))));
},
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 11f38020aa8a1..3b4e789ced691 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -166,7 +166,7 @@ public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucke
ctx,
sub,
starTree,
- MetricStat.SUM.getTypeName(),
+ MetricStat.MIN.getTypeName(),
value -> {
min.set(Math.min(min.get(), (NumericUtils.sortableLongToDouble(value))));
},
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 617ee6a939e6c..c3c7ad299430c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -100,6 +100,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (supportedStarTree != null) {
return getStarTreeCollector(ctx, sub, supportedStarTree);
}
+ System.out.println("nopes nopes");
return getDefaultLeafCollector(ctx, sub);
}
@@ -135,7 +136,7 @@ public void collect(int doc, long bucket) throws IOException {
public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0);
return StarTreeQueryHelper.getStarTreeLeafCollector(
context,
valuesSource,
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index 66b2ba9c98431..ae78e8ca76845 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -36,6 +36,7 @@
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.Query;
+import org.apache.lucene.util.FixedBitSet;
import org.opensearch.action.search.SearchShardTask;
import org.opensearch.action.search.SearchType;
import org.opensearch.common.Nullable;
@@ -130,7 +131,7 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private final List<Releasable> releasables = new CopyOnWriteArrayList<>();
private final AtomicBoolean closed = new AtomicBoolean(false);
private InnerHitsContext innerHitsContext;
- protected volatile Map<LeafReaderContext, StarTreeValuesIterator> starTreeValuesMap;
+ protected volatile Map<LeafReaderContext, FixedBitSet> starTreeValuesMap;
private volatile boolean searchTimedOut;
protected SearchContext() {}
@@ -545,7 +546,7 @@ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryCont
return this;
}
- public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
+ public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
return null;
}
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index fee32ba48e1f5..d56ce506ebddb 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -10,6 +10,7 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
+import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.DocIdSetBuilder;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.index.compositeindex.datacube.Dimension;
@@ -57,48 +58,59 @@ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> pr
* <li>For the remaining columns, use star-tree doc values to match them
* </ul>
*/
- public StarTreeValuesIterator getStarTreeResult() throws IOException {
+ public FixedBitSet getStarTreeResult() throws IOException {
StarTreeResult starTreeResult = traverseStarTree();
- List<StarTreeValuesIterator> andIterators = new ArrayList<>();
- andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator()));
- StarTreeValuesIterator starTreeValuesIterator = andIterators.get(0);
- int length = 0;
- // No matches, return
+ // Initialize FixedBitSet with size maxMatchedDoc + 1
+ FixedBitSet bitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
+ StarTreeValuesIterator starTreeValuesIterator = new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator());
+
+ // No matches, return an empty FixedBitSet
if (starTreeResult.maxMatchedDoc == -1) {
- return starTreeValuesIterator;
+ return bitSet;
+ }
+
+ // Set bits in FixedBitSet for initially matched documents
+ while (starTreeValuesIterator.nextEntry() != NO_MORE_DOCS) {
+ bitSet.set(starTreeValuesIterator.entryId());
}
+
+ // Temporary FixedBitSet reused for filtering
+ FixedBitSet tempBitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
+
+ // Process remaining predicate columns to further filter the results
for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
- DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1);
+
SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) this.starTreeValues.getDimensionValuesIterator(
remainingPredicateColumn
);
- List<Integer> entryIds = new ArrayList<>();
+
long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly
- while (starTreeValuesIterator.nextEntry() != NO_MORE_DOCS) {
- int entryId = starTreeValuesIterator.entryId();
- if (ndv.advance(entryId) > 0) {
+ // Clear the temporary bit set before reuse
+ tempBitSet.clear(0, starTreeResult.maxMatchedDoc + 1);
+
+ // Iterate over the current set of matched document IDs
+ for (int entryId = bitSet.nextSetBit(0); entryId >= 0; entryId = bitSet.nextSetBit(entryId + 1)) {
+ if (ndv.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
final int valuesCount = ndv.valuesCount();
for (int i = 0; i < valuesCount; i++) {
long value = ndv.nextValue();
- // Directly compare value with queryValue
+ // Compare the value with the query value
if (value == queryValue) {
- entryIds.add(entryId);
- break;
+ tempBitSet.set(entryId); // Set bit for the matching entryId
+ break; // No need to check other values for this entryId
}
}
}
}
- DocIdSetBuilder.BulkAdder adder = builder.grow(entryIds.size());
- for (int entryId : entryIds) {
- adder.add(entryId);
- }
- length = entryIds.size();
- starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator());
+
+ // Perform intersection of the current matches with the temp results for this predicate
+ bitSet.and(tempBitSet);
}
- return starTreeValuesIterator;
+
+ return bitSet; // Return the final FixedBitSet with all matches
}
/**
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 2b12c9bba604f..6d184938ccfb5 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -8,10 +8,15 @@
package org.opensearch.search.startree;
+import org.apache.lucene.index.LeafReaderContext;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
+import java.util.ArrayList;
+import java.util.List;
import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
/**
* Query class for querying star tree data structure.
@@ -33,6 +38,18 @@ public class StarTreeQueryContext {
*/
private final Map<String, Long> queryMap;
+// /**
+// * Cache for leaf results
+// * This is used to cache the results for each leaf reader context
+// * to avoid reading the data from the leaf reader context multiple times
+// */
+// private volatile Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> leafResultsCache;
+
+// /**
+// * List of metrics to be computed & cached
+// */
+// private List<StarTreeQueryHelper.MetricInfo> metrics;
+
public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
this.starTree = starTree;
this.queryMap = queryMap;
@@ -45,4 +62,18 @@ public CompositeIndexFieldInfo getStarTree() {
public Map<String, Long> getQueryMap() {
return queryMap;
}
+
+// public void initializeLeafResultsCache() {
+// this.leafResultsCache = new ConcurrentHashMap<>();
+// this.metrics = new ArrayList<>();
+// }
+//
+// public Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> getLeafResultsCache() {
+// return leafResultsCache;
+// }
+//
+// public void addMetric(StarTreeQueryHelper.MetricInfo metric) {
+// metrics.add(metric);
+// }
+
}
From dcb9ac32f335ee00f8f7566c6c6a56065127e490 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 25 Sep 2024 20:05:15 -0700
Subject: [PATCH 22/35] fix initial values in aggs
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 143 ++----------------
.../aggregations/metrics/AvgAggregator.java | 2 +-
.../aggregations/metrics/MaxAggregator.java | 2 +-
.../aggregations/metrics/MinAggregator.java | 2 +-
4 files changed, 13 insertions(+), 136 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index d516a886b6170..179c236775a21 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -214,12 +214,16 @@ public static LeafBucketCollector getStarTreeLeafCollector(
int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
- // Iterate over the FixedBitSet
- for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
- // Advance to the bit (entryId) in the valuesIterator
- if (valuesIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- int count = valuesIterator.valuesCount();
- for (int i = 0; i < count; i++) {
+ if (numBits > 0) {
+ // Iterate over the FixedBitSet
+ for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ // Advance to the bit (entryId) in the valuesIterator
+ if (valuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ continue; // Skip if no more entries
+ }
+
+ // Iterate over the values for the current entryId
+ for (int i = 0, count = valuesIterator.valuesCount(); i < count; i++) {
long value = valuesIterator.nextValue();
valueConsumer.accept(value); // Apply the consumer operation (e.g., max, sum)
}
@@ -237,131 +241,4 @@ public void collect(int doc, long bucket) {
}
};
}
-
-// public static LeafBucketCollector getStarTreeLeafCollectorNew(
-// SearchContext context,
-// ValuesSource.Numeric valuesSource,
-// LeafReaderContext ctx,
-// LeafBucketCollector sub,
-// CompositeIndexFieldInfo starTree,
-// String metric,
-// Consumer<Long> valueConsumer,
-// Runnable finalConsumer
-// ) throws IOException {
-// // Check in contextCache if the star-tree values are already computed
-// Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> cache = context.getStarTreeQueryContext().getLeafResultsCache();
-// if(cache != null) {
-//
-// if (cache.containsKey(ctx)) {
-// MetricInfo metricInfoMap = cache.get(ctx).get(metric);
-// finalConsumer.run();
-// }
-// }
-// else if (!cache.containsKey(ctx)) {
-// // TODO: fetch from cache
-//
-// } else {
-// // TODO: compute cache first
-// }
-//
-// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
-// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
-// String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
-//
-// assert starTreeValues != null;
-// List<SortedNumericStarTreeValuesIterator> valuesIterators;
-// SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
-// metricName
-// );
-// StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
-//
-// int entryId;
-// while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
-// for
-// if (valuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
-// int count = valuesIterator.valuesCount();
-// for (int i = 0; i < count; i++) {
-// long value = valuesIterator.nextValue();
-// valueConsumer.accept(value); // Apply the operation (max, sum, etc.)
-// }
-// }
-// }
-// finalConsumer.run();
-// return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) {
-// @Override
-// public void collect(int doc, long bucket) {
-// throw new CollectionTerminatedException();
-// }
-// };
-// }
-//
-// public abstract class MetricInfo {
-// String metric;
-// MetricStat metricStat;
-//
-//
-//
-// MetricInfo (String metric, MetricStat metricStat) {
-// if (metricStat == MetricStat.SUM) {
-// return new SumMetricInfo(metric);
-// }
-// return null;
-// }
-//
-//
-// public abstract void valueConsumer(long value);
-//
-// public abstract <T extends Number> T getMetricValue();
-// }
-//
-// public class SumMetricInfo extends MetricInfo {
-// CompensatedSum compensatedSum;
-//
-// public SumMetricInfo(String metric) {
-// super(metric, MetricStat.SUM);
-// compensatedSum = new CompensatedSum(0,0);
-// }
-//
-// public void valueConsumer(long value) {
-// compensatedSum.add(NumericUtils.sortableLongToDouble(value));
-// }
-//
-// public Double getMetricValue() {
-// return compensatedSum.value();
-// }
-// }
-//
-// public static void computeLeafResultsCache(SearchContext context,
-// LeafReaderContext ctx,
-// CompositeIndexFieldInfo starTree,
-// List<MetricInfo> metricInfos) throws IOException {
-// Map<String, MetricInfo> leafCache = new HashMap<>();
-// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
-// assert starTreeValues != null;
-// StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues);
-//
-// List<Integer> entryIdCache = new ArrayList<>();
-// int entryId;
-// while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
-// entryIdCache.add(entryId);
-// }
-//
-// for (MetricInfo metricInfo : metricInfos) {
-// SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
-// metricInfo.metric
-// );
-//
-// for (int cachedEntryId : entryIdCache) {
-// if (valuesIterator.advance(cachedEntryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
-// int count = valuesIterator.valuesCount();
-// for (int i = 0; i < count; i++) {
-// long value = valuesIterator.nextValue();
-// metricInfo.valueConsumer(value);
-// }
-// }
-// }
-// leafCache.put(metricInfo.metric, metricInfo);
-// }
-// context.getStarTreeQueryContext().getLeafResultsCache().put(ctx, leafCache);
-// }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index ffd14976b57b0..b375239e94e5b 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -262,7 +262,7 @@ public LeafBucketCollector getStarTreeLeafCollector(
assert starTreeValues != null;
- final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
+ final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0);
SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
sumMetricName
);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
index 338f12ea238e1..257109bca54bb 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
@@ -162,7 +162,7 @@ public void collect(int doc, long bucket) throws IOException {
public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- AtomicReference<Double> max = new AtomicReference<>(Double.NEGATIVE_INFINITY);
+ AtomicReference<Double> max = new AtomicReference<>(maxes.get(0));
return StarTreeQueryHelper.getStarTreeLeafCollector(
context,
valuesSource,
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
index 3b4e789ced691..a9f20bdeb5fd5 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
@@ -159,7 +159,7 @@ public void collect(int doc, long bucket) throws IOException {
public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- AtomicReference<Double> min = new AtomicReference<>(Double.POSITIVE_INFINITY);
+ AtomicReference<Double> min = new AtomicReference<>(mins.get(0));
return StarTreeQueryHelper.getStarTreeLeafCollector(
context,
valuesSource,
From b1bb7165c327d1677b4591b5a9809d1a128e0ea5 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 30 Sep 2024 15:10:11 -0700
Subject: [PATCH 23/35] adding tests
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 42 ++--
.../index/mapper/StarTreeMapper.java | 2 +-
.../search/DefaultSearchContext.java | 28 ---
.../aggregations/metrics/AvgAggregator.java | 144 ++---------
.../aggregations/metrics/SumAggregator.java | 1 -
.../search/internal/SearchContext.java | 20 +-
.../search/startree/StarTreeFilter.java | 1 -
.../search/startree/StarTreeQueryContext.java | 39 +--
.../opensearch/search/SearchServiceTests.java | 159 ++++++-------
.../startree/MetricAggregatorTests.java | 225 ++++++++++++++----
.../aggregations/AggregatorTestCase.java | 138 +++++++----
11 files changed, 418 insertions(+), 381 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 179c236775a21..35b1f1997bcfe 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -12,7 +12,6 @@
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.util.FixedBitSet;
-import org.apache.lucene.util.NumericUtils;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite.CompositeIndexReader;
@@ -30,15 +29,14 @@
import org.opensearch.search.aggregations.AggregatorFactory;
import org.opensearch.search.aggregations.LeafBucketCollector;
import org.opensearch.search.aggregations.LeafBucketCollectorBase;
-import org.opensearch.search.aggregations.metrics.CompensatedSum;
import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory;
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -95,14 +93,16 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
return null;
}
- boolean needCaching = context.aggregations().factories().getFactories().length > 1;
-// List<MetricInfo> metricInfos = new ArrayList<>();
+
for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
- MetricStat metricStat = validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory);
+ MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory);
if (metricStat == null) {
return null;
}
-// metricInfos.add(new )
+ }
+
+ if (context.aggregations().factories().getFactories().length > 1) {
+ context.initializeStarTreeValuesMap();
}
return starTreeQueryContext;
@@ -134,7 +134,7 @@ private static StarTreeQueryContext toStarTreeQueryContext(
/**
* Parse query body to star-tree predicates
- * @param queryBuilder
+ * @param queryBuilder to match supported query shape
* @return predicates to match
*/
private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
@@ -151,11 +151,10 @@ private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder
return predicateMap;
}
- private static MetricStat validateStarTreeMetricSuport(
+ private static MetricStat validateStarTreeMetricSupport(
CompositeDataCubeFieldType compositeIndexFieldInfo,
AggregatorFactory aggregatorFactory
) {
-// List<MetricStat> metricStats = new ArrayList<>();
if (aggregatorFactory instanceof MetricAggregatorFactory && aggregatorFactory.getSubFactories().getFactories().length == 0) {
String field;
Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics()
@@ -197,6 +196,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
Runnable finalConsumer
) throws IOException {
StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
+ assert starTreeValues != null;
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric);
@@ -205,15 +205,10 @@ public static LeafBucketCollector getStarTreeLeafCollector(
metricName
);
// Obtain a FixedBitSet of matched document IDs
- FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues); // Assuming this method gives a FixedBitSet
-
- // Safety check: make sure the FixedBitSet is non-null and valid
- if (matchedDocIds == null) {
- throw new IllegalStateException("FixedBitSet is null");
- }
+ FixedBitSet matchedDocIds = getStarTreeFilteredValues(context, ctx, starTreeValues); // Assuming this method gives a FixedBitSet
+ assert matchedDocIds != null;
int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
-
if (numBits > 0) {
// Iterate over the FixedBitSet
for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
@@ -230,6 +225,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
}
}
+
// Call the final consumer after processing all entries
finalConsumer.run();
@@ -241,4 +237,16 @@ public void collect(int doc, long bucket) {
}
};
}
+
+ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
+ if (context.getStarTreeValuesMap() != null && context.getStarTreeValuesMap().containsKey(ctx)) {
+ return context.getStarTreeValuesMap().get(ctx);
+ }
+
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
+ FixedBitSet result = filter.getStarTreeResult();
+
+ context.getStarTreeValuesMap().put(ctx, result);
+ return result;
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
index 40f05a8b76755..e426e364d42b0 100644
--- a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
+++ b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
@@ -516,7 +516,7 @@ protected void parseCreateField(ParseContext context) {
* @opensearch.experimental
*/
@ExperimentalApi
- public static final class StarTreeFieldType extends CompositeDataCubeFieldType {
+ public static class StarTreeFieldType extends CompositeDataCubeFieldType {
private final StarTreeFieldConfiguration starTreeConfig;
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
index cfe2ad7681a58..feaafd9bbbc6c 100644
--- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -43,7 +43,6 @@
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
-import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.Version;
import org.opensearch.action.search.SearchShardTask;
@@ -59,10 +58,7 @@
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
@@ -185,7 +181,6 @@ final class DefaultSearchContext extends SearchContext {
private SliceBuilder sliceBuilder;
private SearchShardTask task;
private final Version minNodeVersion;
- private StarTreeQueryContext starTreeQueryContext;
/**
* The original query as sent by the user without the types and aliases
@@ -1158,27 +1153,4 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() {
}
return false;
}
-
- @Override
- public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
- this.starTreeQueryContext = starTreeQueryContext;
- return this;
- }
-
- @Override
- public StarTreeQueryContext getStarTreeQueryContext() {
- return this.starTreeQueryContext;
- }
-
- @Override
- public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
- if (this.starTreeValuesMap.containsKey(ctx)) {
- return starTreeValuesMap.get(ctx);
- }
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, this.getStarTreeQueryContext().getQueryMap());
- FixedBitSet result = filter.getStarTreeResult();
-
- starTreeValuesMap.put(ctx, result);
- return result;
- }
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index b375239e94e5b..4a8fe891ae066 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -56,12 +56,11 @@
import org.opensearch.search.aggregations.support.ValuesSource;
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.startree.StarTreeFilter;
import java.io.IOException;
import java.util.Map;
-import java.util.function.Consumer;
+import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeFilteredValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;
/**
@@ -109,13 +108,12 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- // return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
-
final BigArrays bigArrays = context.bigArrays();
final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
@@ -149,105 +147,11 @@ public void collect(int doc, long bucket) throws IOException {
};
}
- public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
- throws IOException {
- final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
- return StarTreeQueryHelper.getStarTreeLeafCollector(
- context,
- valuesSource,
- ctx,
- sub,
- starTree,
- MetricStat.SUM.getTypeName(),
- value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),
- () -> sums.set(0, kahanSummation.value())
- );
- }
-
-// private LeafBucketCollector getStarTreeLeafCollector1(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
-// throws IOException {
-// final BigArrays bigArrays = context.bigArrays();
-// final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
-//
-// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree);
-// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
-// String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
-// starTree.getField(),
-// fieldName,
-// MetricStat.SUM.getTypeName()
-// );
-// assert starTreeValues != null;
-// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName);
-//
-// String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
-// starTree.getField(),
-// fieldName,
-// MetricStat.VALUE_COUNT.getTypeName()
-// );
-// SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName);
-//
-// return new LeafBucketCollectorBase(sub, values) {
-// @Override
-// public void collect(int doc, long bucket) throws IOException {
-// counts = bigArrays.grow(counts, bucket + 1);
-// sums = bigArrays.grow(sums, bucket + 1);
-// compensations = bigArrays.grow(compensations, bucket + 1);
-//
-// if (values.advanceExact(doc) && countValues.advanceExact(doc)) {
-// final long valueCount = values.docValueCount();
-// counts.increment(bucket, countValues.nextValue());
-// // Compute the sum of double values with Kahan summation algorithm which is more
-// // accurate than naive summation.
-// double sum = sums.get(bucket);
-// double compensation = compensations.get(bucket);
-//
-// kahanSummation.reset(sum, compensation);
-//
-// for (int i = 0; i < valueCount; i++) {
-// double value = NumericUtils.sortableLongToDouble(values.nextValue());
-// kahanSummation.add(value);
-// }
-//
-// sums.set(bucket, kahanSummation.value());
-// compensations.set(bucket, kahanSummation.delta());
-// }
-// }
-// };
-// }
-
-
- public LeafBucketCollector getStarTreeCollector2(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
+ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree)
throws IOException {
- final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
- StarTreeQueryHelper.getStarTreeLeafCollector(
- context,
- valuesSource,
- ctx,
- sub,
- starTree,
- MetricStat.SUM.getTypeName(),
- value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)),
- () -> sums.set(0, kahanSummation.value())
- );
- StarTreeQueryHelper.getStarTreeLeafCollector(
- context,
- valuesSource,
- ctx,
- sub,
- starTree,
- MetricStat.VALUE_COUNT.getTypeName(),
- value -> counts.increment(0, value),
- () -> {}
- );
- return LeafBucketCollector.NO_OP_COLLECTOR;
- }
-
- public LeafBucketCollector getStarTreeLeafCollector(
- LeafReaderContext ctx,
- LeafBucketCollector sub,
- CompositeIndexFieldInfo starTree
- ) throws IOException {
StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
+ assert starTreeValues != null;
+
String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName();
String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
starTree.getField(),
@@ -260,35 +164,31 @@ public LeafBucketCollector getStarTreeLeafCollector(
MetricStat.VALUE_COUNT.getTypeName()
);
- assert starTreeValues != null;
final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0);
- SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
- sumMetricName
- );
- SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
- countMetricName
- );
- FixedBitSet matchedDocIds = context.getStarTreeFilteredValues(ctx, starTreeValues);
-
-
- // Safety check: make sure the FixedBitSet is non-null and valid
- if (matchedDocIds == null) {
- throw new IllegalStateException("FixedBitSet is null");
- }
+ SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
+ .getMetricValuesIterator(sumMetricName);
+ SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
+ .getMetricValuesIterator(countMetricName);
+ FixedBitSet matchedDocIds = getStarTreeFilteredValues(context, ctx, starTreeValues);
+ assert matchedDocIds != null;
int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
+ if (numBits > 0) {
+ // Iterate over the FixedBitSet
+ for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ // Advance to the bit (entryId) in the valuesIterator
+ if (sumValuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES
+ || countValueIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ continue; // Skip if no more entries
+ }
- // Iterate over the FixedBitSet
- for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
- // Advance to the bit (entryId) in the valuesIterator
- if (sumValuesIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES &&
- countValueIterator.advance(bit) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- int count = sumValuesIterator.valuesCount();
- for (int i = 0; i < count; i++) {
+ // Iterate over the values for the current entryId
+ for (int i = 0; i < sumValuesIterator.valuesCount(); i++) {
kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue()));
counts.increment(0, countValueIterator.nextValue()); // Apply the consumer operation (e.g., max, sum)
}
+
}
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index c3c7ad299430c..8fb29675eac1c 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -100,7 +100,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
if (supportedStarTree != null) {
return getStarTreeCollector(ctx, sub, supportedStarTree);
}
- System.out.println("nopes nopes");
return getDefaultLeafCollector(ctx, sub);
}
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index ae78e8ca76845..13340fa10f849 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -46,9 +46,7 @@
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.util.BigArrays;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.ObjectMapper;
@@ -81,6 +79,7 @@
import org.opensearch.search.query.ReduceableSearchResult;
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.sort.SortAndFormats;
+import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
@@ -133,6 +132,7 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private InnerHitsContext innerHitsContext;
protected volatile Map<LeafReaderContext, FixedBitSet> starTreeValuesMap;
private volatile boolean searchTimedOut;
+ private StarTreeQueryContext starTreeQueryContext;
protected SearchContext() {}
@@ -538,15 +538,21 @@ public boolean keywordIndexOrDocValuesEnabled() {
return false;
}
- public StarTreeQueryContext getStarTreeQueryContext() {
- return null;
- }
public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
+ this.starTreeQueryContext = starTreeQueryContext;
return this;
}
- public FixedBitSet getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
- return null;
+ public StarTreeQueryContext getStarTreeQueryContext() {
+ return this.starTreeQueryContext;
+ }
+
+ public void initializeStarTreeValuesMap() {
+ this.starTreeValuesMap = new HashMap<>();
+ }
+
+ public Map<LeafReaderContext, FixedBitSet> getStarTreeValuesMap() {
+ return starTreeValuesMap;
}
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index d56ce506ebddb..dcc5bff8a7887 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -10,7 +10,6 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
-import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.DocIdSetBuilder;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.index.compositeindex.datacube.Dimension;
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 6d184938ccfb5..74f05c0b3e7fa 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -8,15 +8,10 @@
package org.opensearch.search.startree;
-import org.apache.lucene.index.LeafReaderContext;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
-import java.util.ArrayList;
-import java.util.List;
import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
/**
* Query class for querying star tree data structure.
@@ -38,17 +33,17 @@ public class StarTreeQueryContext {
*/
private final Map<String, Long> queryMap;
-// /**
-// * Cache for leaf results
-// * This is used to cache the results for each leaf reader context
-// * to avoid reading the data from the leaf reader context multiple times
-// */
-// private volatile Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> leafResultsCache;
+ // /**
+ // * Cache for leaf results
+ // * This is used to cache the results for each leaf reader context
+ // * to avoid reading the data from the leaf reader context multiple times
+ // */
+ // private volatile Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> leafResultsCache;
-// /**
-// * List of metrics to be computed & cached
-// */
-// private List<StarTreeQueryHelper.MetricInfo> metrics;
+ // /**
+ // * List of metrics to be computed & cached
+ // */
+ // private List<StarTreeQueryHelper.MetricInfo> metrics;
public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
this.starTree = starTree;
@@ -62,18 +57,4 @@ public CompositeIndexFieldInfo getStarTree() {
public Map<String, Long> getQueryMap() {
return queryMap;
}
-
-// public void initializeLeafResultsCache() {
-// this.leafResultsCache = new ConcurrentHashMap<>();
-// this.metrics = new ArrayList<>();
-// }
-//
-// public Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> getLeafResultsCache() {
-// return leafResultsCache;
-// }
-//
-// public void addMetric(StarTreeQueryHelper.MetricInfo metric) {
-// metrics.add(metric);
-// }
-
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index 8b8e038a41360..d49111609f95e 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -121,7 +121,6 @@
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.MinAndMax;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.StarTreeQuery;
import org.opensearch.search.suggest.SuggestBuilder;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
import org.opensearch.threadpool.ThreadPool;
@@ -2269,75 +2268,75 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true);
}
- public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
- FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
- setStarTreeIndexSetting("true");
-
- Settings settings = Settings.builder()
- .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
- .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
- .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
- .build();
- CreateIndexRequestBuilder builder = client().admin()
- .indices()
- .prepareCreate("test")
- .setSettings(settings)
- .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
- createIndex("test", builder);
-
- IndicesService indicesService = getInstanceFromNode(IndicesService.class);
- IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
- IndexShard indexShard = indexService.getShard(0);
- ShardSearchRequest request = new ShardSearchRequest(
- OriginalIndices.NONE,
- new SearchRequest().allowPartialSearchResults(true),
- indexShard.shardId(),
- 1,
- new AliasFilter(null, Strings.EMPTY_ARRAY),
- 1.0f,
- -1,
- null,
- null
- );
-
- // Case 1: No query or aggregations, should not use star tree
- SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
- assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-
- // Case 2: MatchAllQuery present but no aggregations, should not use star tree
- sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
- assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-
- // Case 3: MatchAllQuery and aggregations present, should use star tree
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new MatchAllQueryBuilder())
- .aggregation(AggregationBuilders.max("test").field("field"));
- assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-
- // Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new MatchAllQueryBuilder())
- .aggregation(AggregationBuilders.max("test").field("field"))
- .trackTotalHitsUpTo(1000);
- assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-
- // Case 5: TermQuery and aggregations present, should use star tree
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new TermQueryBuilder("sndv", 1))
- .aggregation(AggregationBuilders.max("test").field("field"));
- assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-
- // Case 6: No query, metric aggregations present, should use star tree
- sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
- assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-
- // Case 7: TermQuery and aggregations present, size != 0, should not use star tree
- sourceBuilder = new SearchSourceBuilder().query(new TermQueryBuilder("sndv", 1))
- .aggregation(AggregationBuilders.max("test").field("field"));
- assertQueryType(request, sourceBuilder, IndexOrDocValuesQuery.class);
-
- setStarTreeIndexSetting(null);
- }
+// public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
+// FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+// setStarTreeIndexSetting("true");
+//
+// Settings settings = Settings.builder()
+// .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+// .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+// .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
+// .build();
+// CreateIndexRequestBuilder builder = client().admin()
+// .indices()
+// .prepareCreate("test")
+// .setSettings(settings)
+// .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
+// createIndex("test", builder);
+//
+// IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+// IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
+// IndexShard indexShard = indexService.getShard(0);
+// ShardSearchRequest request = new ShardSearchRequest(
+// OriginalIndices.NONE,
+// new SearchRequest().allowPartialSearchResults(true),
+// indexShard.shardId(),
+// 1,
+// new AliasFilter(null, Strings.EMPTY_ARRAY),
+// 1.0f,
+// -1,
+// null,
+// null
+// );
+//
+// // Case 1: No query or aggregations, should not use star tree
+// SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
+//
+// // Case 2: MatchAllQuery present but no aggregations, should not use star tree
+// sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
+// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
+//
+// // Case 3: MatchAllQuery and aggregations present, should use star tree
+// sourceBuilder = new SearchSourceBuilder().size(0)
+// .query(new MatchAllQueryBuilder())
+// .aggregation(AggregationBuilders.max("test").field("field"));
+// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
+//
+// // Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
+// sourceBuilder = new SearchSourceBuilder().size(0)
+// .query(new MatchAllQueryBuilder())
+// .aggregation(AggregationBuilders.max("test").field("field"))
+// .trackTotalHitsUpTo(1000);
+// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
+//
+// // Case 5: TermQuery and aggregations present, should use star tree
+// sourceBuilder = new SearchSourceBuilder().size(0)
+// .query(new TermQueryBuilder("sndv", 1))
+// .aggregation(AggregationBuilders.max("test").field("field"));
+// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
+//
+// // Case 6: No query, metric aggregations present, should use star tree
+// sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
+// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
+//
+// // Case 7: TermQuery and aggregations present, size != 0, should not use star tree
+// sourceBuilder = new SearchSourceBuilder().query(new TermQueryBuilder("sndv", 1))
+// .aggregation(AggregationBuilders.max("test").field("field"));
+// assertQueryType(request, sourceBuilder, IndexOrDocValuesQuery.class);
+//
+// setStarTreeIndexSetting(null);
+// }
private void setStarTreeIndexSetting(String value) throws IOException {
client().admin()
@@ -2347,14 +2346,14 @@ private void setStarTreeIndexSetting(String value) throws IOException {
.execute();
}
- private void assertQueryType(ShardSearchRequest request, SearchSourceBuilder sourceBuilder, Class<?> expectedQueryClass)
- throws IOException {
- request.source(sourceBuilder);
- SearchService searchService = getInstanceFromNode(SearchService.class);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- assertThat(context.query(), instanceOf(expectedQueryClass));
- searchService.doStop();
- }
- }
+// private void assertQueryType(ShardSearchRequest request, SearchSourceBuilder sourceBuilder, Class<?> expectedQueryClass)
+// throws IOException {
+// request.source(sourceBuilder);
+// SearchService searchService = getInstanceFromNode(SearchService.class);
+// try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+// SearchContext context = searchService.createContext(reader, request, null, true);
+// assertThat(context.query(), instanceOf(expectedQueryClass));
+// searchService.doStop();
+// }
+// }
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 41eba6b548b45..eba7036550987 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -32,9 +32,13 @@
import org.opensearch.index.codec.composite.CompositeIndexReader;
import org.opensearch.index.codec.composite.composite99.Composite99Codec;
import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.NumberFieldMapper;
+import org.opensearch.index.mapper.StarTreeMapper;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.search.aggregations.AggregationBuilder;
import org.opensearch.search.aggregations.AggregatorTestCase;
import org.opensearch.search.aggregations.InternalAggregation;
@@ -48,15 +52,15 @@
import org.opensearch.search.aggregations.metrics.MinAggregationBuilder;
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
-import org.opensearch.search.startree.OriginalOrStarTreeQuery;
-import org.opensearch.search.startree.StarTreeQuery;
import org.junit.After;
import org.junit.Before;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
+import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Function;
@@ -94,6 +98,136 @@ protected Codec getCodec() {
return new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
}
+// public void testStarTreeDocValues1() throws IOException {
+// Directory directory = newDirectory();
+// IndexWriterConfig conf = newIndexWriterConfig(null);
+// conf.setCodec(getCodec());
+// conf.setMergePolicy(newLogMergePolicy());
+// RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
+//
+// Random random = RandomizedTest.getRandom();
+// int totalDocs = 100;
+// final String SNDV = "sndv";
+// final String DV = "dv";
+// int val;
+//
+// // Index 100 random documents
+// for (int i = 0; i < totalDocs; i++) {
+// Document doc = new Document();
+// if (random.nextBoolean()) {
+// val = random.nextInt(10) - 5; // Random long between -5 and 4
+// doc.add(new SortedNumericDocValuesField(SNDV, val));
+// }
+// if (random.nextBoolean()) {
+// val = random.nextInt(20) - 10; // Random long between -10 and 9
+// doc.add(new SortedNumericDocValuesField(DV, val));
+// }
+// if (random.nextBoolean()) {
+// val = random.nextInt(50); // Random long between 0 and 49
+// doc.add(new SortedNumericDocValuesField(FIELD_NAME, val));
+// }
+// iw.addDocument(doc);
+// }
+//
+// if (randomBoolean()) {
+// iw.forceMerge(1);
+// }
+// iw.close();
+//
+// DirectoryReader ir = DirectoryReader.open(directory);
+// initValuesSourceRegistry();
+// LeafReaderContext context = ir.leaves().get(0);
+//
+// SegmentReader reader = Lucene.segmentReader(context.reader());
+// IndexSearcher indexSearcher = newSearcher(reader, false, false);
+// CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+// List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+//
+// CompositeIndexFieldInfo starTree = compositeIndexFields.get(0);
+//
+// SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME);
+// MaxAggregationBuilder maxAggregationBuilder = max("_name").field(FIELD_NAME);
+// MinAggregationBuilder minAggregationBuilder = min("_name").field(FIELD_NAME);
+// ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
+// AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
+//
+// // match-all query
+// Query defaultQuery = new MatchAllDocsQuery();
+// StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates
+// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
+// // numeric-terms query
+// for (int cases = 0; cases < 100; cases++) {
+// Map<String, Long> queryMap;
+// String queryField;
+// long queryValue;
+// if (randomBoolean()) {
+// queryField = SNDV;
+// queryValue = random.nextInt(10);
+// } else {
+// queryField = DV;
+// queryValue = random.nextInt(20) - 15;
+// }
+// defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
+// queryMap = Map.of(queryField, queryValue);
+// starTreeQuery = new StarTreeQuery(starTree, queryMap);
+//
+// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
+// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
+// testCase(
+// indexSearcher,
+// defaultQuery,
+// starTreeQuery,
+// valueCountAggregationBuilder,
+// verifyAggregation(InternalValueCount::getValue)
+// );
+// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
+// }
+// ir.close();
+// directory.close();
+// }
+
+// private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
+// IndexSearcher searcher,
+// Query defaultQuery,
+// T builder,
+// BiConsumer<V, V> verify
+// ) throws IOException {
+//// OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery);
+// V starTreeAggregation = searchAndReduceStarTree(
+// createIndexSettings(),
+// searcher,
+// defaultQuery,
+// builder,
+// DEFAULT_MAX_BUCKETS,
+// false,
+// DEFAULT_MAPPED_FIELD
+// );
+// V expectedAggregation = searchAndReduceStarTree(
+// createIndexSettings(),
+// searcher,
+// defaultQuery,
+// builder,
+// DEFAULT_MAX_BUCKETS,
+// false,
+// DEFAULT_MAPPED_FIELD
+// );
+// verify.accept(expectedAggregation, starTreeAggregation);
+// }
+
+ <T, R extends Number> BiConsumer<T, T> verifyAggregation(Function<T, R> valueExtractor) {
+ return (expectedAggregation, actualAggregation) -> assertEquals(
+ valueExtractor.apply(expectedAggregation).doubleValue(),
+ valueExtractor.apply(actualAggregation).doubleValue(),
+ 0.0f
+ );
+ }
+
+
public void testStarTreeDocValues() throws IOException {
Directory directory = newDirectory();
IndexWriterConfig conf = newIndexWriterConfig(null);
@@ -107,6 +241,7 @@ public void testStarTreeDocValues() throws IOException {
final String DV = "dv";
int val;
+ List<Document> docs = new ArrayList<>();
// Index 100 random documents
for (int i = 0; i < totalDocs; i++) {
Document doc = new Document();
@@ -123,6 +258,7 @@ public void testStarTreeDocValues() throws IOException {
doc.add(new SortedNumericDocValuesField(FIELD_NAME, val));
}
iw.addDocument(doc);
+ docs.add(doc);
}
if (randomBoolean()) {
@@ -137,8 +273,8 @@ public void testStarTreeDocValues() throws IOException {
SegmentReader reader = Lucene.segmentReader(context.reader());
IndexSearcher indexSearcher = newSearcher(reader, false, false);
CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
- List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+ List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
CompositeIndexFieldInfo starTree = compositeIndexFields.get(0);
SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME);
@@ -147,17 +283,17 @@ public void testStarTreeDocValues() throws IOException {
ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
+ Query query = new MatchAllDocsQuery();
// match-all query
- Query defaultQuery = new MatchAllDocsQuery();
- StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates
- testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
- // numeric-terms query
- for (int cases = 0; cases < 100; cases++) {
- Map<String, Long> queryMap;
+ QueryBuilder queryBuilder = null; // no predicates
+// testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
+// testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+// testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+// testCase(indexSearcher, query, queryBuilder, valueCountAggregationBuilder, starTree, verifyAggregation(InternalValueCount::getValue));
+// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
+
+ // Numeric-terms query
+ for (int cases = 0; cases < 1; cases++) {
String queryField;
long queryValue;
if (randomBoolean()) {
@@ -167,39 +303,46 @@ public void testStarTreeDocValues() throws IOException {
queryField = DV;
queryValue = random.nextInt(20) - 15;
}
- defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
- queryMap = Map.of(queryField, queryValue);
- starTreeQuery = new StarTreeQuery(starTree, queryMap);
-
- testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
- testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
- testCase(
- indexSearcher,
- defaultQuery,
- starTreeQuery,
- valueCountAggregationBuilder,
- verifyAggregation(InternalValueCount::getValue)
- );
- testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
+ queryField = DV;
+ queryValue = 1;
+
+ query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
+ queryBuilder = new TermQueryBuilder(queryField, queryValue);
+
+// testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
+// testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+ testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+// testCase(
+// indexSearcher,
+// query,
+// queryBuilder,
+// valueCountAggregationBuilder,
+// starTree,
+// verifyAggregation(InternalValueCount::getValue)
+// );
+// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
}
+
ir.close();
directory.close();
}
+
private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
IndexSearcher searcher,
- Query defaultQuery,
- StarTreeQuery starTreeQuery,
- T builder,
+ Query query,
+ QueryBuilder queryBuilder,
+ T aggBuilder,
+ CompositeIndexFieldInfo starTree,
BiConsumer<V, V> verify
) throws IOException {
- OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery);
V starTreeAggregation = searchAndReduceStarTree(
createIndexSettings(),
searcher,
- originalOrStarTreeQuery,
- builder,
+ query,
+ queryBuilder,
+ aggBuilder,
+ starTree,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
@@ -207,20 +350,14 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
V expectedAggregation = searchAndReduceStarTree(
createIndexSettings(),
searcher,
- defaultQuery,
- builder,
+ query,
+ queryBuilder,
+ aggBuilder,
+ null,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
);
verify.accept(expectedAggregation, starTreeAggregation);
}
-
- <T, R extends Number> BiConsumer<T, T> verifyAggregation(Function<T, R> valueExtractor) {
- return (expectedAggregation, actualAggregation) -> assertEquals(
- valueExtractor.apply(expectedAggregation).doubleValue(),
- valueExtractor.apply(actualAggregation).doubleValue(),
- 0.0f
- );
- }
}
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 63288602abb7e..322bac6cafb2e 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -60,6 +60,7 @@
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.search.AssertingIndexSearcher;
import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.NumericUtils;
import org.opensearch.Version;
import org.opensearch.cluster.metadata.IndexMetadata;
@@ -91,32 +92,19 @@
import org.opensearch.index.cache.bitset.BitsetFilterCache;
import org.opensearch.index.cache.bitset.BitsetFilterCache.Listener;
import org.opensearch.index.cache.query.DisabledQueryCache;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.Dimension;
+import org.opensearch.index.compositeindex.datacube.NumericDimension;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.IndexFieldData;
import org.opensearch.index.fielddata.IndexFieldDataCache;
import org.opensearch.index.fielddata.IndexFieldDataService;
-import org.opensearch.index.mapper.BinaryFieldMapper;
-import org.opensearch.index.mapper.CompletionFieldMapper;
-import org.opensearch.index.mapper.ConstantKeywordFieldMapper;
-import org.opensearch.index.mapper.ContentPath;
-import org.opensearch.index.mapper.DateFieldMapper;
-import org.opensearch.index.mapper.DerivedFieldMapper;
-import org.opensearch.index.mapper.FieldAliasMapper;
-import org.opensearch.index.mapper.FieldMapper;
-import org.opensearch.index.mapper.GeoPointFieldMapper;
-import org.opensearch.index.mapper.GeoShapeFieldMapper;
-import org.opensearch.index.mapper.KeywordFieldMapper;
-import org.opensearch.index.mapper.MappedFieldType;
-import org.opensearch.index.mapper.Mapper;
+import org.opensearch.index.mapper.*;
import org.opensearch.index.mapper.Mapper.BuilderContext;
-import org.opensearch.index.mapper.MapperService;
-import org.opensearch.index.mapper.MatchOnlyTextFieldMapper;
-import org.opensearch.index.mapper.NumberFieldMapper;
-import org.opensearch.index.mapper.ObjectMapper;
import org.opensearch.index.mapper.ObjectMapper.Nested;
-import org.opensearch.index.mapper.RangeFieldMapper;
-import org.opensearch.index.mapper.RangeType;
-import org.opensearch.index.mapper.StarTreeMapper;
-import org.opensearch.index.mapper.TextFieldMapper;
+import org.opensearch.index.query.MatchAllQueryBuilder;
+import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.shard.IndexShard;
import org.opensearch.index.shard.SearchOperationListener;
@@ -135,12 +123,15 @@
import org.opensearch.search.aggregations.support.CoreValuesSourceType;
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
import org.opensearch.search.aggregations.support.ValuesSourceType;
+import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.FetchPhase;
import org.opensearch.search.fetch.subphase.FetchDocValuesPhase;
import org.opensearch.search.fetch.subphase.FetchSourcePhase;
import org.opensearch.search.internal.ContextIndexSearcher;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.lookup.SearchLookup;
+import org.opensearch.search.startree.StarTreeFilter;
+import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.test.InternalAggregationTestCase;
import org.opensearch.test.OpenSearchTestCase;
import org.junit.After;
@@ -152,9 +143,12 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
@@ -331,6 +325,28 @@ protected <A extends Aggregator> A createAggregator(AggregationBuilder aggregati
return aggregator;
}
+ protected CountingAggregator createCountingAggregator(
+ Query query,
+ QueryBuilder queryBuilder,
+ AggregationBuilder aggregationBuilder,
+ IndexSearcher indexSearcher,
+ IndexSettings indexSettings,
+ CompositeIndexFieldInfo starTree,
+ MultiBucketConsumer bucketConsumer,
+ MappedFieldType... fieldTypes
+ ) throws IOException {
+ SearchContext searchContext;
+ if (starTree != null) {
+ searchContext = createSearchContextWithStarTreeContext(indexSearcher, indexSettings, query, queryBuilder, starTree, bucketConsumer, fieldTypes);
+ } else {
+ searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, fieldTypes);
+ }
+ return new CountingAggregator(
+ new AtomicInteger(),
+ createAggregator(aggregationBuilder, searchContext)
+ );
+ }
+
/**
* Create a {@linkplain SearchContext} for testing an {@link Aggregator}.
*/
@@ -344,6 +360,45 @@ protected SearchContext createSearchContext(
return createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
}
+ protected SearchContext createSearchContextWithStarTreeContext(
+ IndexSearcher indexSearcher,
+ IndexSettings indexSettings,
+ Query query,
+ QueryBuilder queryBuilder,
+ CompositeIndexFieldInfo starTree,
+ MultiBucketConsumer bucketConsumer,
+ MappedFieldType... fieldTypes
+ ) throws IOException {
+ SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
+
+// Mock SearchContextAggregations
+ SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class);
+ AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
+ when(searchContext.aggregations()).thenReturn(searchContextAggregations);
+ when(searchContextAggregations.factories()).thenReturn(aggregatorFactories);
+ when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[]{});
+
+ StarTreeMapper.StarTreeFieldType compositeMappedFieldType = mock(StarTreeMapper.StarTreeFieldType.class);
+ when(compositeMappedFieldType.name()).thenReturn(starTree.getField());
+ when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType());
+ Set<CompositeMappedFieldType> compositeFieldTypes = Set.of(compositeMappedFieldType);
+
+ List<Dimension> dimensions = new LinkedList<>();
+ dimensions.add(new NumericDimension("sndv"));
+ dimensions.add(new NumericDimension("dv"));
+ when((compositeMappedFieldType).getDimensions()).thenReturn(dimensions);
+ MapperService mapperService = mock(MapperService.class);
+ when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes);
+ when(searchContext.mapperService()).thenReturn(mapperService);
+
+ SearchSourceBuilder sb = new SearchSourceBuilder().query(queryBuilder);
+ StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(searchContext, sb);
+
+ when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext);
+ return searchContext;
+
+ }
+
protected SearchContext createSearchContext(
IndexSearcher indexSearcher,
IndexSettings indexSettings,
@@ -655,7 +710,9 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
IndexSettings indexSettings,
IndexSearcher searcher,
Query query,
+ QueryBuilder queryBuilder,
AggregationBuilder builder,
+ CompositeIndexFieldInfo compositeIndexFieldInfo,
int maxBucket,
boolean hasNested,
MappedFieldType... fieldTypes
@@ -672,35 +729,14 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
maxBucket,
new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
);
- C root = createAggregator(query, builder, searcher, bucketConsumer, fieldTypes);
-
- if (randomBoolean() && searcher.getIndexReader().leaves().size() > 0) {
- assertTrue(ctx instanceof LeafReaderContext);
- final LeafReaderContext compCTX = (LeafReaderContext) ctx;
- final int size = compCTX.leaves().size();
- final ShardSearcher[] subSearchers = new ShardSearcher[size];
- for (int searcherIDX = 0; searcherIDX < subSearchers.length; searcherIDX++) {
- final LeafReaderContext leave = compCTX.leaves().get(searcherIDX);
- subSearchers[searcherIDX] = new ShardSearcher(leave, compCTX);
- }
- for (ShardSearcher subSearcher : subSearchers) {
- MultiBucketConsumer shardBucketConsumer = new MultiBucketConsumer(
- maxBucket,
- new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
- );
- C a = createAggregator(query, builder, subSearcher, indexSettings, shardBucketConsumer, fieldTypes);
- a.preCollection();
- Weight weight = subSearcher.createWeight(query, ScoreMode.COMPLETE, 1f);
-
- subSearcher.search(weight, a);
- a.postCollection();
- aggs.add(a.buildTopLevel());
- }
- } else {
- root.preCollection();
- searcher.search(query, root);
- root.postCollection();
- aggs.add(root.buildTopLevel());
+ CountingAggregator countingAggregator = createCountingAggregator(query, queryBuilder, builder, searcher, indexSettings, compositeIndexFieldInfo, bucketConsumer, fieldTypes);
+
+ countingAggregator.preCollection();
+ searcher.search(query, countingAggregator);
+ countingAggregator.postCollection();
+ aggs.add(countingAggregator.buildTopLevel());
+ if (compositeIndexFieldInfo != null) {
+ assertEquals(0, countingAggregator.collectCounter.get());
}
MultiBucketConsumer reduceBucketConsumer = new MultiBucketConsumer(
@@ -708,7 +744,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
);
InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forFinalReduction(
- root.context().bigArrays(),
+ countingAggregator.context().bigArrays(),
getMockScriptService(),
reduceBucketConsumer,
pipelines
From c96223c3450031a0de4d07cd6adf6ee7c90a7c62 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 30 Sep 2024 20:14:51 -0700
Subject: [PATCH 24/35] refactoring
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/common/util/FeatureFlags.java | 2 +-
.../startree/utils/StarTreeQueryHelper.java | 42 ++--
.../search/DefaultSearchContext.java | 6 -
.../org/opensearch/search/SearchService.java | 2 +-
.../aggregations/metrics/AvgAggregator.java | 3 +-
.../search/internal/SearchContext.java | 4 -
.../opensearch/search/SearchServiceTests.java | 187 ++++++++++--------
.../startree/MetricAggregatorTests.java | 184 +++--------------
8 files changed, 162 insertions(+), 268 deletions(-)
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index e663d8429da13..6df68013a8119 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 35b1f1997bcfe..ccaaaa597ec0f 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -54,24 +54,19 @@ public class StarTreeQueryHelper {
/**
* Checks if the search context can be supported by star-tree
*/
- public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) {
- boolean canUseStarTree = context.aggregations() != null
- && context.size() == 0
+ public static boolean isStarTreeSupported(SearchContext context) {
+ return context.aggregations() != null
&& context.mapperService().isCompositeIndexPresent()
&& context.parsedPostFilter() == null
&& context.innerHits().getInnerHits().isEmpty()
&& context.sort() == null
- && (!trackTotalHits || context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_DISABLED)
&& context.trackScores() == false
- && context.minimumScore() == null
- && context.terminateAfter() == 0;
- return canUseStarTree;
+ && context.minimumScore() == null;
}
-
/**
- * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
- * Returns null if the query cannot be supported.
+ * Gets StarTreeQueryContext from the search context and source builder.
+ * Returns null if the query & aggregation cannot be supported.
*/
public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
// Current implementation assumes only single star-tree is supported
@@ -93,7 +88,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
return null;
}
-
for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory);
if (metricStat == null) {
@@ -104,20 +98,22 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
if (context.aggregations().factories().getFactories().length > 1) {
context.initializeStarTreeValuesMap();
}
-
return starTreeQueryContext;
}
+ /**
+ * Uses query builder & composite index info to form star-tree query context
+ */
private static StarTreeQueryContext toStarTreeQueryContext(
- CompositeIndexFieldInfo starTree,
- CompositeDataCubeFieldType compositeIndexFieldInfo,
+ CompositeIndexFieldInfo compositeIndexFieldInfo,
+ CompositeDataCubeFieldType compositeFieldType,
QueryBuilder queryBuilder
) {
Map<String, Long> queryMap;
if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
- List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
+ List<String> supportedDimensions = compositeFieldType.getDimensions()
.stream()
.map(Dimension::getField)
.collect(Collectors.toList());
@@ -128,13 +124,12 @@ private static StarTreeQueryContext toStarTreeQueryContext(
} else {
return null;
}
-
- return new StarTreeQueryContext(starTree, queryMap);
+ return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap);
}
/**
* Parse query body to star-tree predicates
- * @param queryBuilder to match supported query shape
+ * @param queryBuilder to match star-tree supported query shape
* @return predicates to match
*/
private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
@@ -185,6 +180,10 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos
return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
}
+ /**
+ * Get the star-tree leaf collector
+ * This collector computes the aggregation prematurely and invokes an early termination collector
+ */
public static LeafBucketCollector getStarTreeLeafCollector(
SearchContext context,
ValuesSource.Numeric valuesSource,
@@ -225,7 +224,6 @@ public static LeafBucketCollector getStarTreeLeafCollector(
}
}
-
// Call the final consumer after processing all entries
finalConsumer.run();
@@ -238,7 +236,11 @@ public void collect(int doc, long bucket) {
};
}
- public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
+ /**
+ * Get the filtered values for the star-tree query
+ */
+ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
+ throws IOException {
if (context.getStarTreeValuesMap() != null && context.getStarTreeValuesMap().containsKey(ctx)) {
return context.getStarTreeValuesMap().get(ctx);
}
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
index feaafd9bbbc6c..74a7482d975df 100644
--- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -34,7 +34,6 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
-import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.BoostQuery;
@@ -43,7 +42,6 @@
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
-import org.apache.lucene.util.FixedBitSet;
import org.opensearch.Version;
import org.opensearch.action.search.SearchShardTask;
import org.opensearch.action.search.SearchType;
@@ -58,7 +56,6 @@
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
@@ -101,8 +98,6 @@
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.slice.SliceBuilder;
import org.opensearch.search.sort.SortAndFormats;
-import org.opensearch.search.startree.StarTreeFilter;
-import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
import java.io.IOException;
@@ -275,7 +270,6 @@ final class DefaultSearchContext extends SearchContext {
this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
- this.starTreeValuesMap = new HashMap<>();
}
@Override
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
index 49f90a9f4c76c..e892a2f1a7620 100644
--- a/server/src/main/java/org/opensearch/search/SearchService.java
+++ b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -1545,7 +1545,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
if (this.indicesService.getCompositeIndexSettings() != null
&& this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
- && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) {
+ && StarTreeQueryHelper.isStarTreeSupported(context)) {
try {
StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source);
if (starTreeQueryContext != null) {
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 4a8fe891ae066..de0fde8df58da 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -108,7 +108,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc
}
CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context);
if (supportedStarTree != null) {
- return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
+ return getStarTreeLeafCollector(ctx, sub, supportedStarTree);
}
return getDefaultLeafCollector(ctx, sub);
}
@@ -164,7 +164,6 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
MetricStat.VALUE_COUNT.getTypeName()
);
-
final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0);
SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
.getMetricValuesIterator(sumMetricName);
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index 13340fa10f849..41eb0b21390f6 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -46,7 +46,6 @@
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.util.BigArrays;
import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.ObjectMapper;
@@ -79,11 +78,9 @@
import org.opensearch.search.query.ReduceableSearchResult;
import org.opensearch.search.rescore.RescoreContext;
import org.opensearch.search.sort.SortAndFormats;
-import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestionSearchContext;
-import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
@@ -538,7 +535,6 @@ public boolean keywordIndexOrDocValuesEnabled() {
return false;
}
-
public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
this.starTreeQueryContext = starTreeQueryContext;
return this;
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index d49111609f95e..b80f1419be761 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -35,8 +35,6 @@
import org.apache.lucene.index.FilterDirectoryReader;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.search.FieldDoc;
-import org.apache.lucene.search.IndexOrDocValuesQuery;
-import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.store.AlreadyClosedException;
import org.opensearch.OpenSearchException;
@@ -78,10 +76,12 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
import org.opensearch.index.compositeindex.CompositeIndexSettings;
import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
import org.opensearch.index.engine.Engine;
+import org.opensearch.index.mapper.CompositeMappedFieldType;
import org.opensearch.index.mapper.DerivedFieldType;
import org.opensearch.index.query.AbstractQueryBuilder;
import org.opensearch.index.query.MatchAllQueryBuilder;
@@ -121,6 +121,7 @@
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.MinAndMax;
import org.opensearch.search.sort.SortOrder;
+import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestBuilder;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
import org.opensearch.threadpool.ThreadPool;
@@ -151,6 +152,7 @@
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -2268,75 +2270,84 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true);
}
-// public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
-// FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
-// setStarTreeIndexSetting("true");
-//
-// Settings settings = Settings.builder()
-// .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
-// .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
-// .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
-// .build();
-// CreateIndexRequestBuilder builder = client().admin()
-// .indices()
-// .prepareCreate("test")
-// .setSettings(settings)
-// .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
-// createIndex("test", builder);
-//
-// IndicesService indicesService = getInstanceFromNode(IndicesService.class);
-// IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
-// IndexShard indexShard = indexService.getShard(0);
-// ShardSearchRequest request = new ShardSearchRequest(
-// OriginalIndices.NONE,
-// new SearchRequest().allowPartialSearchResults(true),
-// indexShard.shardId(),
-// 1,
-// new AliasFilter(null, Strings.EMPTY_ARRAY),
-// 1.0f,
-// -1,
-// null,
-// null
-// );
-//
-// // Case 1: No query or aggregations, should not use star tree
-// SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
-// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-//
-// // Case 2: MatchAllQuery present but no aggregations, should not use star tree
-// sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
-// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-//
-// // Case 3: MatchAllQuery and aggregations present, should use star tree
-// sourceBuilder = new SearchSourceBuilder().size(0)
-// .query(new MatchAllQueryBuilder())
-// .aggregation(AggregationBuilders.max("test").field("field"));
-// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-//
-// // Case 4: MatchAllQuery and aggregations present, but trackTotalHitsUpTo specified, should not use star tree
-// sourceBuilder = new SearchSourceBuilder().size(0)
-// .query(new MatchAllQueryBuilder())
-// .aggregation(AggregationBuilders.max("test").field("field"))
-// .trackTotalHitsUpTo(1000);
-// assertQueryType(request, sourceBuilder, MatchAllDocsQuery.class);
-//
-// // Case 5: TermQuery and aggregations present, should use star tree
-// sourceBuilder = new SearchSourceBuilder().size(0)
-// .query(new TermQueryBuilder("sndv", 1))
-// .aggregation(AggregationBuilders.max("test").field("field"));
-// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-//
-// // Case 6: No query, metric aggregations present, should use star tree
-// sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
-// assertQueryType(request, sourceBuilder, StarTreeQuery.class);
-//
-// // Case 7: TermQuery and aggregations present, size != 0, should not use star tree
-// sourceBuilder = new SearchSourceBuilder().query(new TermQueryBuilder("sndv", 1))
-// .aggregation(AggregationBuilders.max("test").field("field"));
-// assertQueryType(request, sourceBuilder, IndexOrDocValuesQuery.class);
-//
-// setStarTreeIndexSetting(null);
-// }
+ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
+ FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+ setStarTreeIndexSetting("true");
+
+ Settings settings = Settings.builder()
+ .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+ .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+ .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
+ .build();
+ CreateIndexRequestBuilder builder = client().admin()
+ .indices()
+ .prepareCreate("test")
+ .setSettings(settings)
+ .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
+ createIndex("test", builder);
+
+ IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+ IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
+ IndexShard indexShard = indexService.getShard(0);
+ ShardSearchRequest request = new ShardSearchRequest(
+ OriginalIndices.NONE,
+ new SearchRequest().allowPartialSearchResults(true),
+ indexShard.shardId(),
+ 1,
+ new AliasFilter(null, Strings.EMPTY_ARRAY),
+ 1.0f,
+ -1,
+ null,
+ null
+ );
+
+ // Case 1: No query or aggregations, should not use star tree
+ SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 2: MatchAllQuery present but no aggregations, should not use star tree
+ sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 3: MatchAllQuery and aggregations present, should use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ CompositeIndexFieldInfo expectedStarTree = new CompositeIndexFieldInfo(
+ "startree",
+ CompositeMappedFieldType.CompositeFieldType.STAR_TREE
+ );
+ Map<String, Long> expectedQueryMap = null;
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false);
+
+ // Case 4: MatchAllQuery and aggregations present, but trackScores true, should not use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"))
+ .trackScores(true);
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 5: TermQuery and single aggregation, should use star tree, but not initialize query cache
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ expectedQueryMap = Map.of("sndv", 1L);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false);
+
+ // Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"))
+ .aggregation(AggregationBuilders.sum("test2").field("field"));
+ expectedQueryMap = Map.of("sndv", 1L);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), true);
+
+ // Case 7: No query, metric aggregations present, should use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null), false);
+
+ setStarTreeIndexSetting(null);
+ }
private void setStarTreeIndexSetting(String value) throws IOException {
client().admin()
@@ -2346,14 +2357,28 @@ private void setStarTreeIndexSetting(String value) throws IOException {
.execute();
}
-// private void assertQueryType(ShardSearchRequest request, SearchSourceBuilder sourceBuilder, Class<?> expectedQueryClass)
-// throws IOException {
-// request.source(sourceBuilder);
-// SearchService searchService = getInstanceFromNode(SearchService.class);
-// try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
-// SearchContext context = searchService.createContext(reader, request, null, true);
-// assertThat(context.query(), instanceOf(expectedQueryClass));
-// searchService.doStop();
-// }
-// }
+ private void assertStarTreeContext(
+ ShardSearchRequest request,
+ SearchSourceBuilder sourceBuilder,
+ StarTreeQueryContext expectedContext,
+ boolean expectedCacheUsage
+ ) throws IOException {
+ request.source(sourceBuilder);
+ SearchService searchService = getInstanceFromNode(SearchService.class);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ StarTreeQueryContext actualContext = context.getStarTreeQueryContext();
+
+ if (expectedContext == null) {
+ assertThat(context.getStarTreeQueryContext(), nullValue());
+ } else {
+ assertThat(actualContext, notNullValue());
+ assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType());
+ assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField());
+ assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap());
+ }
+ assertThat(context.getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue());
+ searchService.doStop();
+ }
+ }
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index eba7036550987..4d7a8a9f9c644 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -32,11 +32,9 @@
import org.opensearch.index.codec.composite.CompositeIndexReader;
import org.opensearch.index.codec.composite.composite99.Composite99Codec;
import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
-import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.NumberFieldMapper;
-import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.search.aggregations.AggregationBuilder;
@@ -58,9 +56,7 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
-import java.util.Map;
import java.util.Random;
-import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Function;
@@ -98,136 +94,6 @@ protected Codec getCodec() {
return new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
}
-// public void testStarTreeDocValues1() throws IOException {
-// Directory directory = newDirectory();
-// IndexWriterConfig conf = newIndexWriterConfig(null);
-// conf.setCodec(getCodec());
-// conf.setMergePolicy(newLogMergePolicy());
-// RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
-//
-// Random random = RandomizedTest.getRandom();
-// int totalDocs = 100;
-// final String SNDV = "sndv";
-// final String DV = "dv";
-// int val;
-//
-// // Index 100 random documents
-// for (int i = 0; i < totalDocs; i++) {
-// Document doc = new Document();
-// if (random.nextBoolean()) {
-// val = random.nextInt(10) - 5; // Random long between -5 and 4
-// doc.add(new SortedNumericDocValuesField(SNDV, val));
-// }
-// if (random.nextBoolean()) {
-// val = random.nextInt(20) - 10; // Random long between -10 and 9
-// doc.add(new SortedNumericDocValuesField(DV, val));
-// }
-// if (random.nextBoolean()) {
-// val = random.nextInt(50); // Random long between 0 and 49
-// doc.add(new SortedNumericDocValuesField(FIELD_NAME, val));
-// }
-// iw.addDocument(doc);
-// }
-//
-// if (randomBoolean()) {
-// iw.forceMerge(1);
-// }
-// iw.close();
-//
-// DirectoryReader ir = DirectoryReader.open(directory);
-// initValuesSourceRegistry();
-// LeafReaderContext context = ir.leaves().get(0);
-//
-// SegmentReader reader = Lucene.segmentReader(context.reader());
-// IndexSearcher indexSearcher = newSearcher(reader, false, false);
-// CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
-// List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
-//
-// CompositeIndexFieldInfo starTree = compositeIndexFields.get(0);
-//
-// SumAggregationBuilder sumAggregationBuilder = sum("_name").field(FIELD_NAME);
-// MaxAggregationBuilder maxAggregationBuilder = max("_name").field(FIELD_NAME);
-// MinAggregationBuilder minAggregationBuilder = min("_name").field(FIELD_NAME);
-// ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
-// AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
-//
-// // match-all query
-// Query defaultQuery = new MatchAllDocsQuery();
-// StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, null); // no predicates
-// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, valueCountAggregationBuilder, verifyAggregation(InternalValueCount::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
-// // numeric-terms query
-// for (int cases = 0; cases < 100; cases++) {
-// Map<String, Long> queryMap;
-// String queryField;
-// long queryValue;
-// if (randomBoolean()) {
-// queryField = SNDV;
-// queryValue = random.nextInt(10);
-// } else {
-// queryField = DV;
-// queryValue = random.nextInt(20) - 15;
-// }
-// defaultQuery = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
-// queryMap = Map.of(queryField, queryValue);
-// starTreeQuery = new StarTreeQuery(starTree, queryMap);
-//
-// testCase(indexSearcher, defaultQuery, starTreeQuery, sumAggregationBuilder, verifyAggregation(InternalSum::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, maxAggregationBuilder, verifyAggregation(InternalMax::getValue));
-// testCase(indexSearcher, defaultQuery, starTreeQuery, minAggregationBuilder, verifyAggregation(InternalMin::getValue));
-// testCase(
-// indexSearcher,
-// defaultQuery,
-// starTreeQuery,
-// valueCountAggregationBuilder,
-// verifyAggregation(InternalValueCount::getValue)
-// );
-// testCase(indexSearcher, defaultQuery, starTreeQuery, avgAggregationBuilder, verifyAggregation(InternalAvg::getValue));
-// }
-// ir.close();
-// directory.close();
-// }
-
-// private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
-// IndexSearcher searcher,
-// Query defaultQuery,
-// T builder,
-// BiConsumer<V, V> verify
-// ) throws IOException {
-//// OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, defaultQuery);
-// V starTreeAggregation = searchAndReduceStarTree(
-// createIndexSettings(),
-// searcher,
-// defaultQuery,
-// builder,
-// DEFAULT_MAX_BUCKETS,
-// false,
-// DEFAULT_MAPPED_FIELD
-// );
-// V expectedAggregation = searchAndReduceStarTree(
-// createIndexSettings(),
-// searcher,
-// defaultQuery,
-// builder,
-// DEFAULT_MAX_BUCKETS,
-// false,
-// DEFAULT_MAPPED_FIELD
-// );
-// verify.accept(expectedAggregation, starTreeAggregation);
-// }
-
- <T, R extends Number> BiConsumer<T, T> verifyAggregation(Function<T, R> valueExtractor) {
- return (expectedAggregation, actualAggregation) -> assertEquals(
- valueExtractor.apply(expectedAggregation).doubleValue(),
- valueExtractor.apply(actualAggregation).doubleValue(),
- 0.0f
- );
- }
-
-
public void testStarTreeDocValues() throws IOException {
Directory directory = newDirectory();
IndexWriterConfig conf = newIndexWriterConfig(null);
@@ -286,11 +152,18 @@ public void testStarTreeDocValues() throws IOException {
Query query = new MatchAllDocsQuery();
// match-all query
QueryBuilder queryBuilder = null; // no predicates
-// testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
-// testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
-// testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
-// testCase(indexSearcher, query, queryBuilder, valueCountAggregationBuilder, starTree, verifyAggregation(InternalValueCount::getValue));
-// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
+ testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
+ testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+ testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ valueCountAggregationBuilder,
+ starTree,
+ verifyAggregation(InternalValueCount::getValue)
+ );
+ testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
// Numeric-terms query
for (int cases = 0; cases < 1; cases++) {
@@ -303,30 +176,35 @@ public void testStarTreeDocValues() throws IOException {
queryField = DV;
queryValue = random.nextInt(20) - 15;
}
- queryField = DV;
- queryValue = 1;
query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
queryBuilder = new TermQueryBuilder(queryField, queryValue);
-// testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
-// testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
- testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
-// testCase(
-// indexSearcher,
-// query,
-// queryBuilder,
-// valueCountAggregationBuilder,
-// starTree,
-// verifyAggregation(InternalValueCount::getValue)
-// );
-// testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
+ testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
+ // testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+ // testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ valueCountAggregationBuilder,
+ starTree,
+ verifyAggregation(InternalValueCount::getValue)
+ );
+ testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
}
ir.close();
directory.close();
}
+ <T, R extends Number> BiConsumer<T, T> verifyAggregation(Function<T, R> valueExtractor) {
+ return (expectedAggregation, actualAggregation) -> assertEquals(
+ valueExtractor.apply(expectedAggregation).doubleValue(),
+ valueExtractor.apply(actualAggregation).doubleValue(),
+ 0.0f
+ );
+ }
private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
IndexSearcher searcher,
From 7b57f06a2f281a3ff8e43b3d4c3b7e5bf65d4560 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 30 Sep 2024 21:56:20 -0700
Subject: [PATCH 25/35] max/min agg fix, spotless, test fixes
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 16 ++++++++++------
.../SortedNumericStarTreeValuesIterator.java | 4 ++++
.../startree/MetricAggregatorTests.java | 6 +++---
.../search/aggregations/AggregatorTestCase.java | 7 +------
4 files changed, 18 insertions(+), 15 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index ccaaaa597ec0f..a87f6e5f94406 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -20,7 +20,6 @@
import org.opensearch.index.compositeindex.datacube.MetricStat;
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.MatchAllQueryBuilder;
@@ -66,7 +65,7 @@ public static boolean isStarTreeSupported(SearchContext context) {
/**
* Gets StarTreeQueryContext from the search context and source builder.
- * Returns null if the query & aggregation cannot be supported.
+ * Returns null if the query and aggregation cannot be supported.
*/
public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
// Current implementation assumes only single star-tree is supported
@@ -102,7 +101,7 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
}
/**
- * Uses query builder & composite index info to form star-tree query context
+ * Uses query builder and composite index info to form star-tree query context
*/
private static StarTreeQueryContext toStarTreeQueryContext(
CompositeIndexFieldInfo compositeIndexFieldInfo,
@@ -211,8 +210,8 @@ public static LeafBucketCollector getStarTreeLeafCollector(
if (numBits > 0) {
// Iterate over the FixedBitSet
for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
- // Advance to the bit (entryId) in the valuesIterator
- if (valuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ // Advance to the entryId in the valuesIterator
+ if (!valuesIterator.advanceExact(bit)) {
continue; // Skip if no more entries
}
@@ -238,6 +237,8 @@ public void collect(int doc, long bucket) {
/**
* Get the filtered values for the star-tree query
+ * Cache the results in case of multiple aggregations (if cache is initialized)
+ * @return FixedBitSet of matched document IDs
*/
public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
throws IOException {
@@ -248,7 +249,10 @@ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafR
StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
FixedBitSet result = filter.getStarTreeResult();
- context.getStarTreeValuesMap().put(ctx, result);
+ if (context.getStarTreeValuesMap() != null) {
+ context.getStarTreeValuesMap().put(ctx, result);
+ }
return result;
+
}
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
index 44f9545ce4f7f..7d65e9eb1da8e 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
@@ -33,4 +33,8 @@ public long nextValue() throws IOException {
public int valuesCount() throws IOException {
return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
}
+
+ public boolean advanceExact(int target) throws IOException {
+ return ((SortedNumericDocValues) docIdSetIterator).advanceExact(target);
+ }
}
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 4d7a8a9f9c644..06fcae5642e78 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -166,7 +166,7 @@ public void testStarTreeDocValues() throws IOException {
testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
// Numeric-terms query
- for (int cases = 0; cases < 1; cases++) {
+ for (int cases = 0; cases < 100; cases++) {
String queryField;
long queryValue;
if (randomBoolean()) {
@@ -181,8 +181,8 @@ public void testStarTreeDocValues() throws IOException {
queryBuilder = new TermQueryBuilder(queryField, queryValue);
testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
- // testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
- // testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+ testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
+ testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
testCase(
indexSearcher,
query,
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 322bac6cafb2e..e1069bb0ea519 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -60,7 +60,6 @@
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.search.AssertingIndexSearcher;
import org.apache.lucene.util.BytesRef;
-import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.NumericUtils;
import org.opensearch.Version;
import org.opensearch.cluster.metadata.IndexMetadata;
@@ -95,7 +94,6 @@
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.NumericDimension;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.IndexFieldData;
import org.opensearch.index.fielddata.IndexFieldDataCache;
@@ -103,7 +101,6 @@
import org.opensearch.index.mapper.*;
import org.opensearch.index.mapper.Mapper.BuilderContext;
import org.opensearch.index.mapper.ObjectMapper.Nested;
-import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.shard.IndexShard;
@@ -130,7 +127,6 @@
import org.opensearch.search.internal.ContextIndexSearcher;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.search.lookup.SearchLookup;
-import org.opensearch.search.startree.StarTreeFilter;
import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.test.InternalAggregationTestCase;
import org.opensearch.test.OpenSearchTestCase;
@@ -143,7 +139,6 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
-import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
@@ -371,7 +366,7 @@ protected SearchContext createSearchContextWithStarTreeContext(
) throws IOException {
SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
-// Mock SearchContextAggregations
+ // Mock SearchContextAggregations
SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class);
AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
when(searchContext.aggregations()).thenReturn(searchContextAggregations);
From d17db372264c96c30a3e3726b4cf7326a9c8f3ef Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 30 Sep 2024 22:34:45 -0700
Subject: [PATCH 26/35] spotless fix, test refactoring
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/MetricAggregatorTests.java | 92 +++++++++++++++++--
.../aggregations/AggregatorTestCase.java | 73 +++++++++++----
2 files changed, 141 insertions(+), 24 deletions(-)
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 06fcae5642e78..71c07a7c9dd32 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -32,6 +32,8 @@
import org.opensearch.index.codec.composite.CompositeIndexReader;
import org.opensearch.index.codec.composite.composite99.Composite99Codec;
import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.compositeindex.datacube.Dimension;
+import org.opensearch.index.compositeindex.datacube.NumericDimension;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.index.mapper.NumberFieldMapper;
@@ -55,6 +57,7 @@
import java.io.IOException;
import java.util.ArrayList;
+import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.function.BiConsumer;
@@ -149,21 +152,58 @@ public void testStarTreeDocValues() throws IOException {
ValueCountAggregationBuilder valueCountAggregationBuilder = count("_name").field(FIELD_NAME);
AvgAggregationBuilder avgAggregationBuilder = avg("_name").field(FIELD_NAME);
+ List<Dimension> supportedDimensions = new LinkedList<>();
+ supportedDimensions.add(new NumericDimension(SNDV));
+ supportedDimensions.add(new NumericDimension(DV));
+
Query query = new MatchAllDocsQuery();
// match-all query
QueryBuilder queryBuilder = null; // no predicates
- testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
- testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
- testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ sumAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalSum::getValue)
+ );
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ maxAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalMax::getValue)
+ );
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ minAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalMin::getValue)
+ );
testCase(
indexSearcher,
query,
queryBuilder,
valueCountAggregationBuilder,
starTree,
+ supportedDimensions,
verifyAggregation(InternalValueCount::getValue)
);
- testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ avgAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalAvg::getValue)
+ );
// Numeric-terms query
for (int cases = 0; cases < 100; cases++) {
@@ -180,18 +220,51 @@ public void testStarTreeDocValues() throws IOException {
query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue);
queryBuilder = new TermQueryBuilder(queryField, queryValue);
- testCase(indexSearcher, query, queryBuilder, sumAggregationBuilder, starTree, verifyAggregation(InternalSum::getValue));
- testCase(indexSearcher, query, queryBuilder, maxAggregationBuilder, starTree, verifyAggregation(InternalMax::getValue));
- testCase(indexSearcher, query, queryBuilder, minAggregationBuilder, starTree, verifyAggregation(InternalMin::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ sumAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalSum::getValue)
+ );
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ maxAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalMax::getValue)
+ );
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ minAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalMin::getValue)
+ );
testCase(
indexSearcher,
query,
queryBuilder,
valueCountAggregationBuilder,
starTree,
+ supportedDimensions,
verifyAggregation(InternalValueCount::getValue)
);
- testCase(indexSearcher, query, queryBuilder, avgAggregationBuilder, starTree, verifyAggregation(InternalAvg::getValue));
+ testCase(
+ indexSearcher,
+ query,
+ queryBuilder,
+ avgAggregationBuilder,
+ starTree,
+ supportedDimensions,
+ verifyAggregation(InternalAvg::getValue)
+ );
}
ir.close();
@@ -212,6 +285,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
QueryBuilder queryBuilder,
T aggBuilder,
CompositeIndexFieldInfo starTree,
+ List<Dimension> supportedDimensions,
BiConsumer<V, V> verify
) throws IOException {
V starTreeAggregation = searchAndReduceStarTree(
@@ -221,6 +295,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
queryBuilder,
aggBuilder,
starTree,
+ supportedDimensions,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
@@ -232,6 +307,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
queryBuilder,
aggBuilder,
null,
+ null,
DEFAULT_MAX_BUCKETS,
false,
DEFAULT_MAPPED_FIELD
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index e1069bb0ea519..891e2bcfc0df4 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -93,14 +93,34 @@
import org.opensearch.index.cache.query.DisabledQueryCache;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.compositeindex.datacube.Dimension;
-import org.opensearch.index.compositeindex.datacube.NumericDimension;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.fielddata.IndexFieldData;
import org.opensearch.index.fielddata.IndexFieldDataCache;
import org.opensearch.index.fielddata.IndexFieldDataService;
-import org.opensearch.index.mapper.*;
+import org.opensearch.index.mapper.BinaryFieldMapper;
+import org.opensearch.index.mapper.CompletionFieldMapper;
+import org.opensearch.index.mapper.CompositeMappedFieldType;
+import org.opensearch.index.mapper.ConstantKeywordFieldMapper;
+import org.opensearch.index.mapper.ContentPath;
+import org.opensearch.index.mapper.DateFieldMapper;
+import org.opensearch.index.mapper.DerivedFieldMapper;
+import org.opensearch.index.mapper.FieldAliasMapper;
+import org.opensearch.index.mapper.FieldMapper;
+import org.opensearch.index.mapper.GeoPointFieldMapper;
+import org.opensearch.index.mapper.GeoShapeFieldMapper;
+import org.opensearch.index.mapper.KeywordFieldMapper;
+import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.Mapper;
import org.opensearch.index.mapper.Mapper.BuilderContext;
+import org.opensearch.index.mapper.MapperService;
+import org.opensearch.index.mapper.MatchOnlyTextFieldMapper;
+import org.opensearch.index.mapper.NumberFieldMapper;
+import org.opensearch.index.mapper.ObjectMapper;
import org.opensearch.index.mapper.ObjectMapper.Nested;
+import org.opensearch.index.mapper.RangeFieldMapper;
+import org.opensearch.index.mapper.RangeType;
+import org.opensearch.index.mapper.StarTreeMapper;
+import org.opensearch.index.mapper.TextFieldMapper;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.index.shard.IndexShard;
@@ -139,7 +159,6 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
@@ -327,19 +346,26 @@ protected CountingAggregator createCountingAggregator(
IndexSearcher indexSearcher,
IndexSettings indexSettings,
CompositeIndexFieldInfo starTree,
+ List<Dimension> supportedDimensions,
MultiBucketConsumer bucketConsumer,
MappedFieldType... fieldTypes
) throws IOException {
SearchContext searchContext;
if (starTree != null) {
- searchContext = createSearchContextWithStarTreeContext(indexSearcher, indexSettings, query, queryBuilder, starTree, bucketConsumer, fieldTypes);
+ searchContext = createSearchContextWithStarTreeContext(
+ indexSearcher,
+ indexSettings,
+ query,
+ queryBuilder,
+ starTree,
+ supportedDimensions,
+ bucketConsumer,
+ fieldTypes
+ );
} else {
searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, fieldTypes);
}
- return new CountingAggregator(
- new AtomicInteger(),
- createAggregator(aggregationBuilder, searchContext)
- );
+ return new CountingAggregator(new AtomicInteger(), createAggregator(aggregationBuilder, searchContext));
}
/**
@@ -361,27 +387,32 @@ protected SearchContext createSearchContextWithStarTreeContext(
Query query,
QueryBuilder queryBuilder,
CompositeIndexFieldInfo starTree,
+ List<Dimension> supportedDimensions,
MultiBucketConsumer bucketConsumer,
MappedFieldType... fieldTypes
) throws IOException {
- SearchContext searchContext = createSearchContext(indexSearcher, indexSettings, query, bucketConsumer, new NoneCircuitBreakerService(), fieldTypes);
+ SearchContext searchContext = createSearchContext(
+ indexSearcher,
+ indexSettings,
+ query,
+ bucketConsumer,
+ new NoneCircuitBreakerService(),
+ fieldTypes
+ );
// Mock SearchContextAggregations
SearchContextAggregations searchContextAggregations = mock(SearchContextAggregations.class);
AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
when(searchContext.aggregations()).thenReturn(searchContextAggregations);
when(searchContextAggregations.factories()).thenReturn(aggregatorFactories);
- when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[]{});
+ when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {});
StarTreeMapper.StarTreeFieldType compositeMappedFieldType = mock(StarTreeMapper.StarTreeFieldType.class);
when(compositeMappedFieldType.name()).thenReturn(starTree.getField());
when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType());
Set<CompositeMappedFieldType> compositeFieldTypes = Set.of(compositeMappedFieldType);
- List<Dimension> dimensions = new LinkedList<>();
- dimensions.add(new NumericDimension("sndv"));
- dimensions.add(new NumericDimension("dv"));
- when((compositeMappedFieldType).getDimensions()).thenReturn(dimensions);
+ when((compositeMappedFieldType).getDimensions()).thenReturn(supportedDimensions);
MapperService mapperService = mock(MapperService.class);
when(mapperService.getCompositeFieldTypes()).thenReturn(compositeFieldTypes);
when(searchContext.mapperService()).thenReturn(mapperService);
@@ -391,7 +422,6 @@ protected SearchContext createSearchContextWithStarTreeContext(
when(searchContext.getStarTreeQueryContext()).thenReturn(starTreeQueryContext);
return searchContext;
-
}
protected SearchContext createSearchContext(
@@ -708,6 +738,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
QueryBuilder queryBuilder,
AggregationBuilder builder,
CompositeIndexFieldInfo compositeIndexFieldInfo,
+ List<Dimension> supportedDimensions,
int maxBucket,
boolean hasNested,
MappedFieldType... fieldTypes
@@ -724,7 +755,17 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
maxBucket,
new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
);
- CountingAggregator countingAggregator = createCountingAggregator(query, queryBuilder, builder, searcher, indexSettings, compositeIndexFieldInfo, bucketConsumer, fieldTypes);
+ CountingAggregator countingAggregator = createCountingAggregator(
+ query,
+ queryBuilder,
+ builder,
+ searcher,
+ indexSettings,
+ compositeIndexFieldInfo,
+ supportedDimensions,
+ bucketConsumer,
+ fieldTypes
+ );
countingAggregator.preCollection();
searcher.search(query, countingAggregator);
From 6cc085b455bf347aad36d5891c34da1792235816 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 30 Sep 2024 23:04:10 -0700
Subject: [PATCH 27/35] avg aggregator fix, iterators refactoring
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/startree/utils/StarTreeQueryHelper.java | 2 +-
.../startree/utils/iterator/StarTreeValuesIterator.java | 2 +-
.../opensearch/search/aggregations/metrics/AvgAggregator.java | 4 +---
.../java/org/opensearch/search/startree/StarTreeFilter.java | 4 +++-
4 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index a87f6e5f94406..a5bc3327bfb09 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -211,7 +211,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
// Iterate over the FixedBitSet
for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
// Advance to the entryId in the valuesIterator
- if (!valuesIterator.advanceExact(bit)) {
+ if (valuesIterator.advanceExact(bit) == false) {
continue; // Skip if no more entries
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
index 454e5b393973f..32866f3e50092 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
@@ -21,7 +21,7 @@
* @opensearch.experimental
*/
@ExperimentalApi
-public class StarTreeValuesIterator {
+public abstract class StarTreeValuesIterator {
public static final int NO_MORE_ENTRIES = Integer.MAX_VALUE;
protected final DocIdSetIterator docIdSetIterator;
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index de0fde8df58da..49380fdadb175 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -46,7 +46,6 @@
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
-import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
import org.opensearch.index.fielddata.SortedNumericDoubleValues;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.Aggregator;
@@ -177,8 +176,7 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
// Iterate over the FixedBitSet
for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
// Advance to the bit (entryId) in the valuesIterator
- if (sumValuesIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES
- || countValueIterator.advance(bit) == StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ if ((sumValuesIterator.advanceExact(bit) && countValueIterator.advanceExact(bit)) == false) {
continue; // Skip if no more entries
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index dcc5bff8a7887..79a01d5230106 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -62,7 +62,9 @@ public FixedBitSet getStarTreeResult() throws IOException {
// Initialize FixedBitSet with size maxMatchedDoc + 1
FixedBitSet bitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
- StarTreeValuesIterator starTreeValuesIterator = new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator());
+ SortedNumericStarTreeValuesIterator starTreeValuesIterator = new SortedNumericStarTreeValuesIterator(
+ starTreeResult._matchedDocIds.build().iterator()
+ );
// No matches, return an empty FixedBitSet
if (starTreeResult.maxMatchedDoc == -1) {
From 7378fff6f6d4a400c4217b92ceead780a98a87f1 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Wed, 2 Oct 2024 13:36:17 -0700
Subject: [PATCH 28/35] making StarTreeFieldType back to final
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/startree/utils/StarTreeQueryHelper.java | 3 +--
.../main/java/org/opensearch/index/mapper/StarTreeMapper.java | 2 +-
.../opensearch/search/aggregations/metrics/AvgAggregator.java | 1 -
3 files changed, 2 insertions(+), 4 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index a5bc3327bfb09..e4070b96d89c2 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -21,7 +21,6 @@
import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
-import org.opensearch.index.mapper.StarTreeMapper;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.TermQueryBuilder;
@@ -69,7 +68,7 @@ public static boolean isStarTreeSupported(SearchContext context) {
*/
public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
// Current implementation assumes only single star-tree is supported
- CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
+ CompositeDataCubeFieldType compositeMappedFieldType = (CompositeDataCubeFieldType) context.mapperService()
.getCompositeFieldTypes()
.iterator()
.next();
diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
index e426e364d42b0..40f05a8b76755 100644
--- a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
+++ b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
@@ -516,7 +516,7 @@ protected void parseCreateField(ParseContext context) {
* @opensearch.experimental
*/
@ExperimentalApi
- public static class StarTreeFieldType extends CompositeDataCubeFieldType {
+ public static final class StarTreeFieldType extends CompositeDataCubeFieldType {
private final StarTreeFieldConfiguration starTreeConfig;
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 49380fdadb175..4299ebc475058 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -185,7 +185,6 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue()));
counts.increment(0, countValueIterator.nextValue()); // Apply the consumer operation (e.g., max, sum)
}
-
}
}
From a58e3dfbeec5f045a9ba00d0af926ba9c4868109 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Sun, 6 Oct 2024 20:54:36 -0700
Subject: [PATCH 29/35] move value cache to star tree context + other comments
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 51 ++----
.../search/internal/SearchContext.java | 11 --
.../search/startree/StarTreeQueryContext.java | 31 ++--
.../search/SearchServiceStarTreeTests.java | 156 ++++++++++++++++++
.../opensearch/search/SearchServiceTests.java | 122 --------------
.../aggregations/AggregatorTestCase.java | 3 +-
6 files changed, 194 insertions(+), 180 deletions(-)
create mode 100644 server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index e4070b96d89c2..53c5e2bfd13e4 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -53,13 +53,7 @@ public class StarTreeQueryHelper {
* Checks if the search context can be supported by star-tree
*/
public static boolean isStarTreeSupported(SearchContext context) {
- return context.aggregations() != null
- && context.mapperService().isCompositeIndexPresent()
- && context.parsedPostFilter() == null
- && context.innerHits().getInnerHits().isEmpty()
- && context.sort() == null
- && context.trackScores() == false
- && context.minimumScore() == null;
+ return context.aggregations() != null && context.mapperService().isCompositeIndexPresent() && context.parsedPostFilter() == null;
}
/**
@@ -77,15 +71,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
compositeMappedFieldType.getCompositeIndexType()
);
- StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.toStarTreeQueryContext(
- starTree,
- compositeMappedFieldType,
- source.query()
- );
- if (starTreeQueryContext == null) {
- return null;
- }
-
for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory);
if (metricStat == null) {
@@ -93,10 +78,9 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
}
}
- if (context.aggregations().factories().getFactories().length > 1) {
- context.initializeStarTreeValuesMap();
- }
- return starTreeQueryContext;
+ boolean cacheStarTreeValues = context.aggregations().factories().getFactories().length > 1;
+
+ return StarTreeQueryHelper.toStarTreeQueryContext(starTree, compositeMappedFieldType, source.query(), cacheStarTreeValues);
}
/**
@@ -105,7 +89,8 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
private static StarTreeQueryContext toStarTreeQueryContext(
CompositeIndexFieldInfo compositeIndexFieldInfo,
CompositeDataCubeFieldType compositeFieldType,
- QueryBuilder queryBuilder
+ QueryBuilder queryBuilder,
+ boolean cacheStarTreeValues
) {
Map<String, Long> queryMap;
if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
@@ -122,7 +107,7 @@ private static StarTreeQueryContext toStarTreeQueryContext(
} else {
return null;
}
- return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap);
+ return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap, cacheStarTreeValues);
}
/**
@@ -201,14 +186,14 @@ public static LeafBucketCollector getStarTreeLeafCollector(
SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
metricName
);
- // Obtain a FixedBitSet of matched document IDs
- FixedBitSet matchedDocIds = getStarTreeFilteredValues(context, ctx, starTreeValues); // Assuming this method gives a FixedBitSet
- assert matchedDocIds != null;
+ // Obtain a FixedBitSet of matched star tree document IDs
+ FixedBitSet filteredValues = getStarTreeFilteredValues(context, ctx, starTreeValues);
+ assert filteredValues != null;
- int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
+ int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs)
if (numBits > 0) {
- // Iterate over the FixedBitSet
- for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ // Iterate over the filtered values
+ for (int bit = filteredValues.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? filteredValues.nextSetBit(bit + 1) : -1) {
// Advance to the entryId in the valuesIterator
if (valuesIterator.advanceExact(bit) == false) {
continue; // Skip if no more entries
@@ -241,17 +226,17 @@ public void collect(int doc, long bucket) {
*/
public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
throws IOException {
- if (context.getStarTreeValuesMap() != null && context.getStarTreeValuesMap().containsKey(ctx)) {
- return context.getStarTreeValuesMap().get(ctx);
+ Map<LeafReaderContext, FixedBitSet> valueCache = context.getStarTreeQueryContext().getStarTreeValuesMap();
+ if (valueCache != null && valueCache.containsKey(ctx)) {
+ return valueCache.get(ctx);
}
StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
FixedBitSet result = filter.getStarTreeResult();
- if (context.getStarTreeValuesMap() != null) {
- context.getStarTreeValuesMap().put(ctx, result);
+ if (valueCache != null) {
+ valueCache.put(ctx, result);
}
return result;
-
}
}
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
index 41eb0b21390f6..b7ea06d2989e5 100644
--- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -31,12 +31,10 @@
package org.opensearch.search.internal;
-import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.Query;
-import org.apache.lucene.util.FixedBitSet;
import org.opensearch.action.search.SearchShardTask;
import org.opensearch.action.search.SearchType;
import org.opensearch.common.Nullable;
@@ -127,7 +125,6 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private final List<Releasable> releasables = new CopyOnWriteArrayList<>();
private final AtomicBoolean closed = new AtomicBoolean(false);
private InnerHitsContext innerHitsContext;
- protected volatile Map<LeafReaderContext, FixedBitSet> starTreeValuesMap;
private volatile boolean searchTimedOut;
private StarTreeQueryContext starTreeQueryContext;
@@ -543,12 +540,4 @@ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryCont
public StarTreeQueryContext getStarTreeQueryContext() {
return this.starTreeQueryContext;
}
-
- public void initializeStarTreeValuesMap() {
- this.starTreeValuesMap = new HashMap<>();
- }
-
- public Map<LeafReaderContext, FixedBitSet> getStarTreeValuesMap() {
- return starTreeValuesMap;
- }
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 74f05c0b3e7fa..39fb3d3a68dda 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -8,10 +8,13 @@
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;
import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
/**
* Query class for querying star tree data structure.
@@ -31,23 +34,21 @@ public class StarTreeQueryContext {
* Map of field name to a value to be queried for that field
* This is used to filter the data based on the query
*/
- private final Map<String, Long> queryMap;
+ private volatile Map<String, Long> queryMap;
- // /**
- // * Cache for leaf results
- // * This is used to cache the results for each leaf reader context
- // * to avoid reading the data from the leaf reader context multiple times
- // */
- // private volatile Map<LeafReaderContext, Map<String, StarTreeQueryHelper.MetricInfo>> leafResultsCache;
-
- // /**
- // * List of metrics to be computed & cached
- // */
- // private List<StarTreeQueryHelper.MetricInfo> metrics;
+ /**
+ * Cache for leaf results
+ * This is used to cache the results for each leaf reader context
+ * to avoid reading the filtered values from the leaf reader context multiple times
+ */
+ protected volatile Map<LeafReaderContext, FixedBitSet> starTreeValuesMap;
- public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap) {
+ public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, boolean cacheStarTreeValues) {
this.starTree = starTree;
this.queryMap = queryMap;
+ if (cacheStarTreeValues) {
+ starTreeValuesMap = new ConcurrentHashMap<>();
+ }
}
public CompositeIndexFieldInfo getStarTree() {
@@ -57,4 +58,8 @@ public CompositeIndexFieldInfo getStarTree() {
public Map<String, Long> getQueryMap() {
return queryMap;
}
+
+ public Map<LeafReaderContext, FixedBitSet> getStarTreeValuesMap() {
+ return starTreeValuesMap;
+ }
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
new file mode 100644
index 0000000000000..b2953f5482c9e
--- /dev/null
+++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
@@ -0,0 +1,156 @@
+/*
+ * 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;
+
+import org.opensearch.action.OriginalIndices;
+import org.opensearch.action.admin.indices.create.CreateIndexRequestBuilder;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
+import org.opensearch.core.common.Strings;
+import org.opensearch.index.IndexService;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.compositeindex.CompositeIndexSettings;
+import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
+import org.opensearch.index.mapper.CompositeMappedFieldType;
+import org.opensearch.index.query.MatchAllQueryBuilder;
+import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.index.shard.IndexShard;
+import org.opensearch.indices.IndicesService;
+import org.opensearch.search.aggregations.AggregationBuilders;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.internal.AliasFilter;
+import org.opensearch.search.internal.ReaderContext;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.internal.ShardSearchRequest;
+import org.opensearch.search.startree.StarTreeQueryContext;
+import org.opensearch.test.OpenSearchSingleNodeTestCase;
+
+import java.io.IOException;
+import java.util.Map;
+
+import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
+
+public class SearchServiceStarTreeTests extends OpenSearchSingleNodeTestCase {
+
+ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
+ FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+ setStarTreeIndexSetting("true");
+
+ Settings settings = Settings.builder()
+ .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+ .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+ .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
+ .build();
+ CreateIndexRequestBuilder builder = client().admin()
+ .indices()
+ .prepareCreate("test")
+ .setSettings(settings)
+ .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
+ createIndex("test", builder);
+
+ IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+ IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
+ IndexShard indexShard = indexService.getShard(0);
+ ShardSearchRequest request = new ShardSearchRequest(
+ OriginalIndices.NONE,
+ new SearchRequest().allowPartialSearchResults(true),
+ indexShard.shardId(),
+ 1,
+ new AliasFilter(null, Strings.EMPTY_ARRAY),
+ 1.0f,
+ -1,
+ null,
+ null
+ );
+
+ // Case 1: No query or aggregations, should not use star tree
+ SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 2: MatchAllQuery present but no aggregations, should not use star tree
+ sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 3: MatchAllQuery and aggregations present, should use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ CompositeIndexFieldInfo expectedStarTree = new CompositeIndexFieldInfo(
+ "startree",
+ CompositeMappedFieldType.CompositeFieldType.STAR_TREE
+ );
+ Map<String, Long> expectedQueryMap = null;
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, false), false);
+
+ // Case 4: MatchAllQuery and aggregations present, but postFilter specified, should not use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new MatchAllQueryBuilder())
+ .aggregation(AggregationBuilders.max("test").field("field"))
+ .postFilter(new MatchAllQueryBuilder());
+ assertStarTreeContext(request, sourceBuilder, null, false);
+
+ // Case 5: TermQuery and single aggregation, should use star tree, but not initialize query cache
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"));
+ expectedQueryMap = Map.of("sndv", 1L);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, false), false);
+
+ // Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache
+ sourceBuilder = new SearchSourceBuilder().size(0)
+ .query(new TermQueryBuilder("sndv", 1))
+ .aggregation(AggregationBuilders.max("test").field("field"))
+ .aggregation(AggregationBuilders.sum("test2").field("field"));
+ expectedQueryMap = Map.of("sndv", 1L);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, true), true);
+
+ // Case 7: No query, metric aggregations present, should use star tree
+ sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null, false), false);
+
+ setStarTreeIndexSetting(null);
+ }
+
+ private void setStarTreeIndexSetting(String value) throws IOException {
+ client().admin()
+ .cluster()
+ .prepareUpdateSettings()
+ .setTransientSettings(Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), value).build())
+ .execute();
+ }
+
+ private void assertStarTreeContext(
+ ShardSearchRequest request,
+ SearchSourceBuilder sourceBuilder,
+ StarTreeQueryContext expectedContext,
+ boolean expectedCacheUsage
+ ) throws IOException {
+ request.source(sourceBuilder);
+ SearchService searchService = getInstanceFromNode(SearchService.class);
+ try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
+ SearchContext context = searchService.createContext(reader, request, null, true);
+ StarTreeQueryContext actualContext = context.getStarTreeQueryContext();
+
+ if (expectedContext == null) {
+ assertThat(context.getStarTreeQueryContext(), nullValue());
+ } else {
+ assertThat(actualContext, notNullValue());
+ assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType());
+ assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField());
+ assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap());
+ assertThat(context.getStarTreeQueryContext().getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue());
+ }
+ searchService.doStop();
+ }
+ }
+}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceTests.java b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
index b80f1419be761..514e99a126267 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceTests.java
@@ -39,7 +39,6 @@
import org.apache.lucene.store.AlreadyClosedException;
import org.opensearch.OpenSearchException;
import org.opensearch.action.OriginalIndices;
-import org.opensearch.action.admin.indices.create.CreateIndexRequestBuilder;
import org.opensearch.action.index.IndexResponse;
import org.opensearch.action.search.ClearScrollRequest;
import org.opensearch.action.search.DeletePitResponse;
@@ -55,14 +54,12 @@
import org.opensearch.action.support.IndicesOptions;
import org.opensearch.action.support.PlainActionFuture;
import org.opensearch.action.support.WriteRequest;
-import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.UUIDs;
import org.opensearch.common.settings.Setting;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.settings.SettingsException;
import org.opensearch.common.unit.TimeValue;
-import org.opensearch.common.util.FeatureFlags;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.io.stream.StreamInput;
@@ -76,12 +73,7 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.index.IndexService;
import org.opensearch.index.IndexSettings;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
-import org.opensearch.index.compositeindex.CompositeIndexSettings;
-import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
import org.opensearch.index.engine.Engine;
-import org.opensearch.index.mapper.CompositeMappedFieldType;
import org.opensearch.index.mapper.DerivedFieldType;
import org.opensearch.index.query.AbstractQueryBuilder;
import org.opensearch.index.query.MatchAllQueryBuilder;
@@ -121,7 +113,6 @@
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.MinAndMax;
import org.opensearch.search.sort.SortOrder;
-import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.search.suggest.SuggestBuilder;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
import org.opensearch.threadpool.ThreadPool;
@@ -152,7 +143,6 @@
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.notNullValue;
-import static org.hamcrest.CoreMatchers.nullValue;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -2269,116 +2259,4 @@ public void testCanMatchSearchAfterDescLessThanMinWithTrackTotalhits() throws IO
primarySort.order(SortOrder.DESC);
assertEquals(SearchService.canMatchSearchAfter(searchAfter, minMax, primarySort, 1000), true);
}
-
- public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
- FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
- setStarTreeIndexSetting("true");
-
- Settings settings = Settings.builder()
- .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
- .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
- .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true)
- .build();
- CreateIndexRequestBuilder builder = client().admin()
- .indices()
- .prepareCreate("test")
- .setSettings(settings)
- .setMapping(StarTreeDocValuesFormatTests.getExpandedMapping());
- createIndex("test", builder);
-
- IndicesService indicesService = getInstanceFromNode(IndicesService.class);
- IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test"));
- IndexShard indexShard = indexService.getShard(0);
- ShardSearchRequest request = new ShardSearchRequest(
- OriginalIndices.NONE,
- new SearchRequest().allowPartialSearchResults(true),
- indexShard.shardId(),
- 1,
- new AliasFilter(null, Strings.EMPTY_ARRAY),
- 1.0f,
- -1,
- null,
- null
- );
-
- // Case 1: No query or aggregations, should not use star tree
- SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
- assertStarTreeContext(request, sourceBuilder, null, false);
-
- // Case 2: MatchAllQuery present but no aggregations, should not use star tree
- sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
- assertStarTreeContext(request, sourceBuilder, null, false);
-
- // Case 3: MatchAllQuery and aggregations present, should use star tree
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new MatchAllQueryBuilder())
- .aggregation(AggregationBuilders.max("test").field("field"));
- CompositeIndexFieldInfo expectedStarTree = new CompositeIndexFieldInfo(
- "startree",
- CompositeMappedFieldType.CompositeFieldType.STAR_TREE
- );
- Map<String, Long> expectedQueryMap = null;
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false);
-
- // Case 4: MatchAllQuery and aggregations present, but trackScores true, should not use star tree
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new MatchAllQueryBuilder())
- .aggregation(AggregationBuilders.max("test").field("field"))
- .trackScores(true);
- assertStarTreeContext(request, sourceBuilder, null, false);
-
- // Case 5: TermQuery and single aggregation, should use star tree, but not initialize query cache
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new TermQueryBuilder("sndv", 1))
- .aggregation(AggregationBuilders.max("test").field("field"));
- expectedQueryMap = Map.of("sndv", 1L);
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), false);
-
- // Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache
- sourceBuilder = new SearchSourceBuilder().size(0)
- .query(new TermQueryBuilder("sndv", 1))
- .aggregation(AggregationBuilders.max("test").field("field"))
- .aggregation(AggregationBuilders.sum("test2").field("field"));
- expectedQueryMap = Map.of("sndv", 1L);
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap), true);
-
- // Case 7: No query, metric aggregations present, should use star tree
- sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null), false);
-
- setStarTreeIndexSetting(null);
- }
-
- private void setStarTreeIndexSetting(String value) throws IOException {
- client().admin()
- .cluster()
- .prepareUpdateSettings()
- .setTransientSettings(Settings.builder().put(CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING.getKey(), value).build())
- .execute();
- }
-
- private void assertStarTreeContext(
- ShardSearchRequest request,
- SearchSourceBuilder sourceBuilder,
- StarTreeQueryContext expectedContext,
- boolean expectedCacheUsage
- ) throws IOException {
- request.source(sourceBuilder);
- SearchService searchService = getInstanceFromNode(SearchService.class);
- try (ReaderContext reader = searchService.createOrGetReaderContext(request, false)) {
- SearchContext context = searchService.createContext(reader, request, null, true);
- StarTreeQueryContext actualContext = context.getStarTreeQueryContext();
-
- if (expectedContext == null) {
- assertThat(context.getStarTreeQueryContext(), nullValue());
- } else {
- assertThat(actualContext, notNullValue());
- assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType());
- assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField());
- assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap());
- }
- assertThat(context.getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue());
- searchService.doStop();
- }
- }
}
diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
index 891e2bcfc0df4..e1728c4476699 100644
--- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java
@@ -99,6 +99,7 @@
import org.opensearch.index.fielddata.IndexFieldDataService;
import org.opensearch.index.mapper.BinaryFieldMapper;
import org.opensearch.index.mapper.CompletionFieldMapper;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
import org.opensearch.index.mapper.CompositeMappedFieldType;
import org.opensearch.index.mapper.ConstantKeywordFieldMapper;
import org.opensearch.index.mapper.ContentPath;
@@ -407,7 +408,7 @@ protected SearchContext createSearchContextWithStarTreeContext(
when(searchContextAggregations.factories()).thenReturn(aggregatorFactories);
when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[] {});
- StarTreeMapper.StarTreeFieldType compositeMappedFieldType = mock(StarTreeMapper.StarTreeFieldType.class);
+ CompositeDataCubeFieldType compositeMappedFieldType = mock(CompositeDataCubeFieldType.class);
when(compositeMappedFieldType.name()).thenReturn(starTree.getField());
when(compositeMappedFieldType.getCompositeIndexType()).thenReturn(starTree.getType());
Set<CompositeMappedFieldType> compositeFieldTypes = Set.of(compositeMappedFieldType);
From 47c39d88af3812707a155a827c4fa812805704d6 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Tue, 8 Oct 2024 18:47:06 -0700
Subject: [PATCH 30/35] refactor cache map to cache array
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../startree/utils/StarTreeQueryHelper.java | 30 +++++++++----------
.../aggregations/metrics/AvgAggregator.java | 5 +++-
.../aggregations/metrics/SumAggregator.java | 10 +++----
.../search/startree/StarTreeQueryContext.java | 29 +++++++++++++-----
.../search/SearchServiceStarTreeTests.java | 22 ++++++++------
5 files changed, 58 insertions(+), 38 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 53c5e2bfd13e4..3fe982fc0091a 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -11,6 +11,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.CollectionTerminatedException;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
@@ -78,19 +79,21 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
}
}
+ // need to cache star tree values only for multiple aggregations
boolean cacheStarTreeValues = context.aggregations().factories().getFactories().length > 1;
+ int cacheSize = cacheStarTreeValues ? context.indexShard().segments(false).size() : -1;
- return StarTreeQueryHelper.toStarTreeQueryContext(starTree, compositeMappedFieldType, source.query(), cacheStarTreeValues);
+ return StarTreeQueryHelper.tryCreateStarTreeQueryContext(starTree, compositeMappedFieldType, source.query(), cacheSize);
}
/**
* Uses query builder and composite index info to form star-tree query context
*/
- private static StarTreeQueryContext toStarTreeQueryContext(
+ private static StarTreeQueryContext tryCreateStarTreeQueryContext(
CompositeIndexFieldInfo compositeIndexFieldInfo,
CompositeDataCubeFieldType compositeFieldType,
QueryBuilder queryBuilder,
- boolean cacheStarTreeValues
+ int cacheStarTreeValuesSize
) {
Map<String, Long> queryMap;
if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
@@ -107,7 +110,7 @@ private static StarTreeQueryContext toStarTreeQueryContext(
} else {
return null;
}
- return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap, cacheStarTreeValues);
+ return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap, cacheStarTreeValuesSize);
}
/**
@@ -193,7 +196,9 @@ public static LeafBucketCollector getStarTreeLeafCollector(
int numBits = filteredValues.length(); // Get the number of the filtered values (matching docs)
if (numBits > 0) {
// Iterate over the filtered values
- for (int bit = filteredValues.nextSetBit(0); bit != -1; bit = (bit + 1 < numBits) ? filteredValues.nextSetBit(bit + 1) : -1) {
+ for (int bit = filteredValues.nextSetBit(0); bit != DocIdSetIterator.NO_MORE_DOCS; bit = (bit + 1 < numBits)
+ ? filteredValues.nextSetBit(bit + 1)
+ : DocIdSetIterator.NO_MORE_DOCS) {
// Advance to the entryId in the valuesIterator
if (valuesIterator.advanceExact(bit) == false) {
continue; // Skip if no more entries
@@ -226,16 +231,11 @@ public void collect(int doc, long bucket) {
*/
public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
throws IOException {
- Map<LeafReaderContext, FixedBitSet> valueCache = context.getStarTreeQueryContext().getStarTreeValuesMap();
- if (valueCache != null && valueCache.containsKey(ctx)) {
- return valueCache.get(ctx);
- }
-
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
- FixedBitSet result = filter.getStarTreeResult();
-
- if (valueCache != null) {
- valueCache.put(ctx, result);
+ FixedBitSet result = context.getStarTreeQueryContext().getStarTreeValues(ctx);
+ if (result == null) {
+ StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
+ result = filter.getStarTreeResult();
+ context.getStarTreeQueryContext().setStarTreeValues(ctx, result);
}
return result;
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 4299ebc475058..0a3a514701e15 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -33,6 +33,7 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.CollectionTerminatedException;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.NumericUtils;
@@ -174,7 +175,9 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
int numBits = matchedDocIds.length(); // Get the length of the FixedBitSet
if (numBits > 0) {
// Iterate over the FixedBitSet
- for (int bit = matchedDocIds.nextSetBit(0); bit != -1; bit = bit + 1 < numBits ? matchedDocIds.nextSetBit(bit + 1) : -1) {
+ for (int bit = matchedDocIds.nextSetBit(0); bit != DocIdSetIterator.NO_MORE_DOCS; bit = bit + 1 < numBits
+ ? matchedDocIds.nextSetBit(bit + 1)
+ : DocIdSetIterator.NO_MORE_DOCS) {
// Advance to the bit (entryId) in the valuesIterator
if ((sumValuesIterator.advanceExact(bit) && countValueIterator.advanceExact(bit)) == false) {
continue; // Skip if no more entries
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
index 8fb29675eac1c..3d237a94c5699 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -62,13 +62,13 @@
*/
public class SumAggregator extends NumericMetricsAggregator.SingleValue {
- protected final ValuesSource.Numeric valuesSource;
- protected final DocValueFormat format;
+ private final ValuesSource.Numeric valuesSource;
+ private final DocValueFormat format;
- protected DoubleArray sums;
- protected DoubleArray compensations;
+ private DoubleArray sums;
+ private DoubleArray compensations;
- public SumAggregator(
+ SumAggregator(
String name,
ValuesSourceConfig valuesSourceConfig,
SearchContext context,
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 39fb3d3a68dda..165de139ec0d4 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -14,7 +14,6 @@
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
/**
* Query class for querying star tree data structure.
@@ -34,20 +33,20 @@ public class StarTreeQueryContext {
* Map of field name to a value to be queried for that field
* This is used to filter the data based on the query
*/
- private volatile Map<String, Long> queryMap;
+ private final Map<String, Long> queryMap;
/**
* Cache for leaf results
* This is used to cache the results for each leaf reader context
* to avoid reading the filtered values from the leaf reader context multiple times
*/
- protected volatile Map<LeafReaderContext, FixedBitSet> starTreeValuesMap;
+ private FixedBitSet[] starTreeValues;
- public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, boolean cacheStarTreeValues) {
+ public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, int cacheStarTreeValuesSize) {
this.starTree = starTree;
this.queryMap = queryMap;
- if (cacheStarTreeValues) {
- starTreeValuesMap = new ConcurrentHashMap<>();
+ if (cacheStarTreeValuesSize > -1) {
+ starTreeValues = new FixedBitSet[cacheStarTreeValuesSize];
}
}
@@ -59,7 +58,21 @@ public Map<String, Long> getQueryMap() {
return queryMap;
}
- public Map<LeafReaderContext, FixedBitSet> getStarTreeValuesMap() {
- return starTreeValuesMap;
+ public FixedBitSet[] getStarTreeValues() {
+ return starTreeValues;
}
+
+ public FixedBitSet getStarTreeValues(LeafReaderContext ctx) {
+ if (starTreeValues != null) {
+ return starTreeValues[ctx.ord];
+ }
+ return null;
+ }
+
+ public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) {
+ if (starTreeValues != null) {
+ starTreeValues[ctx.ord] = values;
+ }
+ }
+
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
index b2953f5482c9e..88d5265dfda78 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
@@ -75,11 +75,11 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
// Case 1: No query or aggregations, should not use star tree
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
- assertStarTreeContext(request, sourceBuilder, null, false);
+ assertStarTreeContext(request, sourceBuilder, null, -1);
// Case 2: MatchAllQuery present but no aggregations, should not use star tree
sourceBuilder = new SearchSourceBuilder().query(new MatchAllQueryBuilder());
- assertStarTreeContext(request, sourceBuilder, null, false);
+ assertStarTreeContext(request, sourceBuilder, null, -1);
// Case 3: MatchAllQuery and aggregations present, should use star tree
sourceBuilder = new SearchSourceBuilder().size(0)
@@ -90,21 +90,21 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
CompositeMappedFieldType.CompositeFieldType.STAR_TREE
);
Map<String, Long> expectedQueryMap = null;
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, false), false);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, -1), -1);
// Case 4: MatchAllQuery and aggregations present, but postFilter specified, should not use star tree
sourceBuilder = new SearchSourceBuilder().size(0)
.query(new MatchAllQueryBuilder())
.aggregation(AggregationBuilders.max("test").field("field"))
.postFilter(new MatchAllQueryBuilder());
- assertStarTreeContext(request, sourceBuilder, null, false);
+ assertStarTreeContext(request, sourceBuilder, null, -1);
// Case 5: TermQuery and single aggregation, should use star tree, but not initialize query cache
sourceBuilder = new SearchSourceBuilder().size(0)
.query(new TermQueryBuilder("sndv", 1))
.aggregation(AggregationBuilders.max("test").field("field"));
expectedQueryMap = Map.of("sndv", 1L);
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, false), false);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, -1), -1);
// Case 6: TermQuery and multiple aggregations present, should use star tree & initialize cache
sourceBuilder = new SearchSourceBuilder().size(0)
@@ -112,11 +112,11 @@ public void testParseQueryToOriginalOrStarTreeQuery() throws IOException {
.aggregation(AggregationBuilders.max("test").field("field"))
.aggregation(AggregationBuilders.sum("test2").field("field"));
expectedQueryMap = Map.of("sndv", 1L);
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, true), true);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, expectedQueryMap, 0), 0);
// Case 7: No query, metric aggregations present, should use star tree
sourceBuilder = new SearchSourceBuilder().size(0).aggregation(AggregationBuilders.max("test").field("field"));
- assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null, false), false);
+ assertStarTreeContext(request, sourceBuilder, new StarTreeQueryContext(expectedStarTree, null, -1), -1);
setStarTreeIndexSetting(null);
}
@@ -133,7 +133,7 @@ private void assertStarTreeContext(
ShardSearchRequest request,
SearchSourceBuilder sourceBuilder,
StarTreeQueryContext expectedContext,
- boolean expectedCacheUsage
+ int expectedCacheUsage
) throws IOException {
request.source(sourceBuilder);
SearchService searchService = getInstanceFromNode(SearchService.class);
@@ -148,7 +148,11 @@ private void assertStarTreeContext(
assertEquals(expectedContext.getStarTree().getType(), actualContext.getStarTree().getType());
assertEquals(expectedContext.getStarTree().getField(), actualContext.getStarTree().getField());
assertEquals(expectedContext.getQueryMap(), actualContext.getQueryMap());
- assertThat(context.getStarTreeQueryContext().getStarTreeValuesMap(), expectedCacheUsage ? notNullValue() : nullValue());
+ if (expectedCacheUsage > -1) {
+ assertEquals(expectedCacheUsage, actualContext.getStarTreeValues().length);
+ } else {
+ assertNull(actualContext.getStarTreeValues());
+ }
}
searchService.doStop();
}
From 5101f532f63e8d821b38d9a31a469f7582e335ce Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Thu, 10 Oct 2024 11:50:48 -0700
Subject: [PATCH 31/35] refactor star tree filter
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../datacube/DateDimension.java | 5 ++
.../compositeindex/datacube/Dimension.java | 3 ++
.../datacube/NumericDimension.java | 6 +++
.../datacube/ReadDimension.java | 6 +++
.../startree/utils/StarTreeQueryHelper.java | 10 +++-
.../search/startree/StarTreeFilter.java | 53 +++++++++----------
.../search/startree/StarTreeQueryContext.java | 10 ++--
.../search/SearchServiceStarTreeTests.java | 2 +-
.../startree/MetricAggregatorTests.java | 8 +--
9 files changed, 63 insertions(+), 40 deletions(-)
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java
index ee6d5b4680c73..8feb9ccd27dbd 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java
@@ -8,6 +8,7 @@
package org.opensearch.index.compositeindex.datacube;
+import org.apache.lucene.index.DocValuesType;
import org.opensearch.common.Rounding;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.common.time.DateUtils;
@@ -169,4 +170,8 @@ public int compare(DateTimeUnitRounding unit1, DateTimeUnitRounding unit2) {
public static List<DateTimeUnitRounding> getSortedDateTimeUnits(List<DateTimeUnitRounding> dateTimeUnits) {
return dateTimeUnits.stream().sorted(new DateTimeUnitComparator()).collect(Collectors.toList());
}
+
+ public DocValuesType getDocValuesType() {
+ return DocValuesType.SORTED_NUMERIC;
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java
index cfa8d3a2a8164..3d71b38881693 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java
@@ -8,6 +8,7 @@
package org.opensearch.index.compositeindex.datacube;
+import org.apache.lucene.index.DocValuesType;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.core.xcontent.ToXContent;
@@ -42,4 +43,6 @@ public interface Dimension extends ToXContent {
* Returns the list of dimension fields that represent the dimension
*/
List<String> getSubDimensionNames();
+
+ DocValuesType getDocValuesType();
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java
index acc14f5f05c68..f1d1b15337f4a 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java
@@ -8,6 +8,7 @@
package org.opensearch.index.compositeindex.datacube;
+import org.apache.lucene.index.DocValuesType;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
@@ -71,4 +72,9 @@ public boolean equals(Object o) {
public int hashCode() {
return Objects.hash(field);
}
+
+ @Override
+ public DocValuesType getDocValuesType() {
+ return DocValuesType.SORTED_NUMERIC;
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java
index be3667f10b6da..0e2ec086abc0a 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java
@@ -8,6 +8,7 @@
package org.opensearch.index.compositeindex.datacube;
+import org.apache.lucene.index.DocValuesType;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.index.mapper.CompositeDataCubeFieldType;
@@ -69,4 +70,9 @@ public boolean equals(Object o) {
public int hashCode() {
return Objects.hash(field);
}
+
+ @Override
+ public DocValuesType getDocValuesType() {
+ return DocValuesType.SORTED_NUMERIC;
+ }
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index 3fe982fc0091a..b362b083a8d20 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -8,6 +8,7 @@
package org.opensearch.index.compositeindex.datacube.startree.utils;
+import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.CollectionTerminatedException;
@@ -99,6 +100,12 @@ private static StarTreeQueryContext tryCreateStarTreeQueryContext(
if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
+ // TODO: Add support for keyword fields
+ if (compositeFieldType.getDimensions().stream().anyMatch(d -> d.getDocValuesType() != DocValuesType.SORTED_NUMERIC)) {
+ // return null for non-numeric fields
+ return null;
+ }
+
List<String> supportedDimensions = compositeFieldType.getDimensions()
.stream()
.map(Dimension::getField)
@@ -233,8 +240,7 @@ public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafR
throws IOException {
FixedBitSet result = context.getStarTreeQueryContext().getStarTreeValues(ctx);
if (result == null) {
- StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
- result = filter.getStarTreeResult();
+ result = StarTreeFilter.getStarTreeResult(starTreeValues, context.getStarTreeQueryContext().getQueryMap());
context.getStarTreeQueryContext().setStarTreeValues(ctx, result);
}
return result;
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index 79a01d5230106..b0b60956b551f 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -42,28 +42,18 @@
public class StarTreeFilter {
private static final Logger logger = LogManager.getLogger(StarTreeFilter.class);
- private final Map<String, Long> queryMap;
- private final StarTreeValues starTreeValues;
-
- public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map<String, Long> predicateEvaluators) {
- // This filter operator does not support AND/OR/NOT operations as of now.
- starTreeValues = starTreeAggrStructure;
- queryMap = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
- }
-
/**
- * <ul>
- * <li>First go over the star tree and try to match as many dimensions as possible
- * <li>For the remaining columns, use star-tree doc values to match them
- * </ul>
+ * First go over the star tree and try to match as many dimensions as possible
+ * For the remaining columns, use star-tree doc values to match them
*/
- public FixedBitSet getStarTreeResult() throws IOException {
- StarTreeResult starTreeResult = traverseStarTree();
+ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, Map<String, Long> predicateEvaluators) throws IOException {
+ Map<String, Long> queryMap = predicateEvaluators != null ? predicateEvaluators : Collections.emptyMap();
+ StarTreeResult starTreeResult = traverseStarTree(starTreeValues, queryMap);
// Initialize FixedBitSet with size maxMatchedDoc + 1
FixedBitSet bitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
SortedNumericStarTreeValuesIterator starTreeValuesIterator = new SortedNumericStarTreeValuesIterator(
- starTreeResult._matchedDocIds.build().iterator()
+ starTreeResult.matchedDocIds.build().iterator()
);
// No matches, return an empty FixedBitSet
@@ -80,10 +70,10 @@ public FixedBitSet getStarTreeResult() throws IOException {
FixedBitSet tempBitSet = new FixedBitSet(starTreeResult.maxMatchedDoc + 1);
// Process remaining predicate columns to further filter the results
- for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) {
+ for (String remainingPredicateColumn : starTreeResult.remainingPredicateColumns) {
logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc);
- SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) this.starTreeValues.getDimensionValuesIterator(
+ SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) starTreeValues.getDimensionValuesIterator(
remainingPredicateColumn
);
@@ -118,8 +108,8 @@ public FixedBitSet getStarTreeResult() throws IOException {
* Helper method to traverse the star tree, get matching documents and keep track of all the
* predicate dimensions that are not matched.
*/
- private StarTreeResult traverseStarTree() throws IOException {
- DocIdSetBuilder docsWithField = new DocIdSetBuilder(this.starTreeValues.getStarTreeDocumentCount());
+ private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, Map<String, Long> queryMap) throws IOException {
+ DocIdSetBuilder docsWithField = new DocIdSetBuilder(starTreeValues.getStarTreeDocumentCount());
DocIdSetBuilder.BulkAdder adder;
Set<String> globalRemainingPredicateColumns = null;
StarTreeNode starTree = starTreeValues.getRoot();
@@ -214,15 +204,20 @@ private StarTreeResult traverseStarTree() throws IOException {
/**
* Helper class to wrap the result from traversing the star tree.
* */
- static class StarTreeResult {
- final DocIdSetBuilder _matchedDocIds;
- final Set<String> _remainingPredicateColumns;
- final int numOfMatchedDocs;
- final int maxMatchedDoc;
-
- StarTreeResult(DocIdSetBuilder matchedDocIds, Set<String> remainingPredicateColumns, int numOfMatchedDocs, int maxMatchedDoc) {
- _matchedDocIds = matchedDocIds;
- _remainingPredicateColumns = remainingPredicateColumns;
+ public static class StarTreeResult {
+ public final DocIdSetBuilder matchedDocIds;
+ public final Set<String> remainingPredicateColumns;
+ public final int numOfMatchedDocs;
+ public final int maxMatchedDoc;
+
+ public StarTreeResult(
+ DocIdSetBuilder matchedDocIds,
+ Set<String> remainingPredicateColumns,
+ int numOfMatchedDocs,
+ int maxMatchedDoc
+ ) {
+ this.matchedDocIds = matchedDocIds;
+ this.remainingPredicateColumns = remainingPredicateColumns;
this.numOfMatchedDocs = numOfMatchedDocs;
this.maxMatchedDoc = maxMatchedDoc;
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 165de139ec0d4..5a53dd122cdb0 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -40,13 +40,15 @@ public class StarTreeQueryContext {
* This is used to cache the results for each leaf reader context
* to avoid reading the filtered values from the leaf reader context multiple times
*/
- private FixedBitSet[] starTreeValues;
+ private final FixedBitSet[] starTreeValues;
- public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, int cacheStarTreeValuesSize) {
+ public StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map<String, Long> queryMap, int numSegmentsCache) {
this.starTree = starTree;
this.queryMap = queryMap;
- if (cacheStarTreeValuesSize > -1) {
- starTreeValues = new FixedBitSet[cacheStarTreeValuesSize];
+ if (numSegmentsCache > -1) {
+ starTreeValues = new FixedBitSet[numSegmentsCache];
+ } else {
+ starTreeValues = null;
}
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
index 88d5265dfda78..0c88154ca2b38 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
@@ -17,7 +17,7 @@
import org.opensearch.core.common.Strings;
import org.opensearch.index.IndexService;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests;
import org.opensearch.index.compositeindex.CompositeIndexSettings;
import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
import org.opensearch.index.mapper.CompositeMappedFieldType;
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
index 71c07a7c9dd32..0327bd9990784 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -13,7 +13,7 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.codecs.Codec;
-import org.apache.lucene.codecs.lucene99.Lucene99Codec;
+import org.apache.lucene.codecs.lucene912.Lucene912Codec;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.SortedNumericDocValuesField;
import org.apache.lucene.index.DirectoryReader;
@@ -30,8 +30,8 @@
import org.opensearch.common.util.FeatureFlags;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.index.codec.composite.CompositeIndexReader;
-import org.opensearch.index.codec.composite.composite99.Composite99Codec;
-import org.opensearch.index.codec.composite99.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.codec.composite.composite912.Composite912Codec;
+import org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests;
import org.opensearch.index.compositeindex.datacube.Dimension;
import org.opensearch.index.compositeindex.datacube.NumericDimension;
import org.opensearch.index.mapper.MappedFieldType;
@@ -94,7 +94,7 @@ protected Codec getCodec() {
} catch (IOException e) {
throw new RuntimeException(e);
}
- return new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger);
+ return new Composite912Codec(Lucene912Codec.Mode.BEST_SPEED, mapperService, testLogger);
}
public void testStarTreeDocValues() throws IOException {
From 1c74afbd5db276ad159ef440ad812c6904cfdbac Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 14 Oct 2024 23:37:22 -0700
Subject: [PATCH 32/35] Rebasing & adding star tree filter tests
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../search/startree/StarTreeFilter.java | 33 +-
.../StarTreeDocValuesFormatTests.java | 2 +-
.../startree/StarTreeFilterTests.java | 319 ++++++++++++++++++
3 files changed, 338 insertions(+), 16 deletions(-)
create mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index b0b60956b551f..f1cc1f11a671a 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -10,12 +10,13 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
+import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.DocIdSetBuilder;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.index.compositeindex.datacube.Dimension;
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.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
@@ -82,16 +83,20 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, Map<S
// Clear the temporary bit set before reuse
tempBitSet.clear(0, starTreeResult.maxMatchedDoc + 1);
- // Iterate over the current set of matched document IDs
- for (int entryId = bitSet.nextSetBit(0); entryId >= 0; entryId = bitSet.nextSetBit(entryId + 1)) {
- if (ndv.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- final int valuesCount = ndv.valuesCount();
- for (int i = 0; i < valuesCount; i++) {
- long value = ndv.nextValue();
- // Compare the value with the query value
- if (value == queryValue) {
- tempBitSet.set(entryId); // Set bit for the matching entryId
- break; // No need to check other values for this entryId
+ if (bitSet.length() > 0) {
+ // Iterate over the current set of matched document IDs
+ for (int entryId = bitSet.nextSetBit(0); entryId != DocIdSetIterator.NO_MORE_DOCS; entryId = (entryId + 1 < bitSet.length())
+ ? bitSet.nextSetBit(entryId + 1)
+ : DocIdSetIterator.NO_MORE_DOCS) {
+ if (ndv.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
+ final int valuesCount = ndv.valuesCount();
+ for (int i = 0; i < valuesCount; i++) {
+ long value = ndv.nextValue();
+ // Compare the value with the query value
+ if (value == queryValue) {
+ tempBitSet.set(entryId); // Set bit for the matching entryId
+ break; // No need to check other values for this entryId
+ }
}
}
}
@@ -119,13 +124,11 @@ private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, Ma
.map(Dimension::getField)
.collect(Collectors.toList());
boolean foundLeafNode = starTree.isLeaf();
+ assert foundLeafNode == false; // root node is never leaf
Queue<StarTreeNode> queue = new ArrayDeque<>();
queue.add(starTree);
int currentDimensionId = -1;
Set<String> remainingPredicateColumns = new HashSet<>(queryMap.keySet());
- if (foundLeafNode) {
- globalRemainingPredicateColumns = new HashSet<>(remainingPredicateColumns);
- }
int matchedDocsCountInStarTree = 0;
int maxDocNum = -1;
StarTreeNode starTreeNode;
@@ -180,7 +183,7 @@ private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, Ma
Iterator<? extends StarTreeNode> childrenIterator = starTreeNode.getChildrenIterator();
while (childrenIterator.hasNext()) {
StarTreeNode childNode = childrenIterator.next();
- if (childNode.getDimensionValue() != StarTreeUtils.ALL) {
+ if (childNode.getStarTreeNodeType() != StarTreeNodeType.STAR.getValue()) {
queue.add(childNode);
foundLeafNode |= childNode.isLeaf();
}
diff --git a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
index 42ee5a5336dd1..f081cadc1362c 100644
--- a/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
+++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java
@@ -361,7 +361,7 @@ public static XContentBuilder getExpandedMapping() throws IOException {
});
}
- private static XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
+ public static XContentBuilder topMapping(CheckedConsumer<XContentBuilder, IOException> buildFields) throws IOException {
XContentBuilder builder = XContentFactory.jsonBuilder().startObject().startObject("_doc");
buildFields.accept(builder);
return builder.endObject().endObject();
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
new file mode 100644
index 0000000000000..009184d24817d
--- /dev/null
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
@@ -0,0 +1,319 @@
+/*
+ * 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.aggregations.startree;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.lucene912.Lucene912Codec;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.IndexableField;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.util.FixedBitSet;
+import org.opensearch.common.lucene.Lucene;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
+import org.opensearch.index.codec.composite.composite912.Composite912Codec;
+import org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper;
+import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.mapper.MapperService;
+import org.opensearch.search.aggregations.AggregatorTestCase;
+import org.opensearch.search.startree.StarTreeFilter;
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests.topMapping;
+
+public class StarTreeFilterTests extends AggregatorTestCase {
+
+ private static final String FIELD_NAME = "field";
+ private static final String SNDV = "sndv";
+ private static final String SDV = "sdv";
+ private static final String DV = "dv";
+
+ @Before
+ public void setup() {
+ FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build());
+ }
+
+ @After
+ public void teardown() throws IOException {
+ FeatureFlags.initializeFeatureFlags(Settings.EMPTY);
+ }
+
+ protected Codec getCodec(int maxLeafDoc, boolean skipStarNodeCreationForSDVDimension) {
+ final Logger testLogger = LogManager.getLogger(StarTreeFilterTests.class);
+ MapperService mapperService;
+ try {
+ mapperService = StarTreeDocValuesFormatTests.createMapperService(
+ getExpandedMapping(maxLeafDoc, skipStarNodeCreationForSDVDimension)
+ );
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ return new Composite912Codec(Lucene912Codec.Mode.BEST_SPEED, mapperService, testLogger);
+ }
+
+ public void testStarTreeFilterWithNoDocsInSVDField() throws IOException {
+ testStarTreeFilter(5, true);
+ }
+
+ public void testStarTreeFilterWithDocsInSVDFieldButNoStarNode() throws IOException {
+ testStarTreeFilter(10, false);
+ }
+
+ private void testStarTreeFilter(int maxLeafDoc, boolean skipStarNodeCreationForSDVDimension) throws IOException {
+ Directory directory = newDirectory();
+ IndexWriterConfig conf = newIndexWriterConfig(null);
+ conf.setCodec(getCodec(maxLeafDoc, skipStarNodeCreationForSDVDimension));
+ conf.setMergePolicy(newLogMergePolicy());
+ RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf);
+ int totalDocs = 100;
+
+ List<Document> docs = new ArrayList<>();
+ for (int i = 0; i < totalDocs; i++) {
+ Document doc = new Document();
+ doc.add(new SortedNumericDocValuesField(SNDV, i));
+ doc.add(new SortedNumericDocValuesField(DV, 2 * i));
+ doc.add(new SortedNumericDocValuesField(FIELD_NAME, 3 * i));
+ if (skipStarNodeCreationForSDVDimension) {
+ // adding SDV field only star node creation is skipped for SDV dimension
+ doc.add(new SortedNumericDocValuesField(SDV, 4 * i));
+ }
+ iw.addDocument(doc);
+ docs.add(doc);
+ }
+ iw.forceMerge(1);
+ iw.close();
+
+ DirectoryReader ir = DirectoryReader.open(directory);
+ initValuesSourceRegistry();
+ LeafReaderContext context = ir.leaves().get(0);
+ SegmentReader reader = Lucene.segmentReader(context.reader());
+ CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+
+ long starTreeDocCount, docCount;
+
+ // assert that all documents are included if no filters are given
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(), context);
+ docCount = getDocCount(docs, Map.of());
+ assertEquals(totalDocs, starTreeDocCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // single filter - matches docs
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SNDV, 0L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, 0L));
+ assertEquals(1, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // single filter on 3rd field in ordered dimension - matches docs
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(DV, 0L), context);
+ docCount = getDocCount(docs, Map.of(DV, 0L));
+ assertEquals(1, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // single filter - does not match docs
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SNDV, 101L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, 101L));
+ assertEquals(0, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // single filter on 3rd field in ordered dimension - does not match docs
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(DV, -101L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, -101L));
+ assertEquals(0, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // multiple filters - matches docs
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SNDV, 0L, DV, 0L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, 0L, DV, 0L));
+ assertEquals(1, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // no document should match the filter
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SNDV, 0L, DV, -11L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, 0L, DV, -11L));
+ assertEquals(0, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // Only the first filter should match some documents, second filter matches none
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SNDV, 0L, DV, -100L), context);
+ docCount = getDocCount(docs, Map.of(SNDV, 0L, DV, -100L));
+ assertEquals(0, docCount);
+ assertEquals(docCount, starTreeDocCount);
+
+ // non-dimension fields in filter - should throw IllegalArgumentException
+ expectThrows(
+ IllegalArgumentException.class,
+ () -> getDocCountFromStarTree(starTreeDocValuesReader, Map.of(FIELD_NAME, 0L), context)
+ );
+
+ if (skipStarNodeCreationForSDVDimension == true) {
+ // Documents are not indexed
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SDV, 4L), context);
+ docCount = getDocCount(docs, Map.of(SDV, 4L));
+ assertEquals(1, docCount);
+ assertEquals(docCount, starTreeDocCount);
+ } else {
+ // Documents are indexed
+ starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, Map.of(SDV, 4L), context);
+ docCount = getDocCount(docs, Map.of(SDV, 4L));
+ assertEquals(0, docCount);
+ assertEquals(docCount, starTreeDocCount);
+ }
+
+ ir.close();
+ directory.close();
+ }
+
+ // Counts the documents having field SNDV & applied filters
+ private long getDocCount(List<Document> documents, Map<String, Long> filters) {
+ long count = 0;
+ for (Document doc : documents) {
+ // Check if SNDV field is present
+ IndexableField sndvField = doc.getField(SNDV);
+ if (sndvField == null) continue; // Skip if SNDV is not present
+
+ // Apply filters if provided
+ if (!filters.isEmpty()) {
+ boolean matches = filters.entrySet().stream().allMatch(entry -> {
+ IndexableField field = doc.getField(entry.getKey());
+ return field != null && field.numericValue().longValue() == entry.getValue();
+ });
+ if (!matches) continue;
+ }
+
+ // Increment count if the document passes all conditions
+ count++;
+ }
+ return count;
+ }
+
+ // Returns count of documents in the star tree having field SNDV & applied filters
+ private long getDocCountFromStarTree(CompositeIndexReader starTreeDocValuesReader, Map<String, Long> filters, LeafReaderContext context)
+ throws IOException {
+ List<CompositeIndexFieldInfo> compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields();
+ CompositeIndexFieldInfo starTree = compositeIndexFields.get(0);
+ StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(context, starTree);
+ FixedBitSet filteredValues = StarTreeFilter.getStarTreeResult(starTreeValues, filters);
+
+ SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(
+ StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(
+ starTree.getField(),
+ SNDV,
+ MetricStat.VALUE_COUNT.getTypeName()
+ )
+ );
+
+ long docCount = 0;
+ int numBits = filteredValues.length();
+ if (numBits > 0) {
+ for (int bit = filteredValues.nextSetBit(0); bit != DocIdSetIterator.NO_MORE_DOCS; bit = (bit + 1 < numBits)
+ ? filteredValues.nextSetBit(bit + 1)
+ : DocIdSetIterator.NO_MORE_DOCS) {
+
+ // Assert that we can advance to the document ID in the values iterator
+ boolean canAdvance = valuesIterator.advanceExact(bit);
+ assert canAdvance : "Cannot advance to document ID " + bit + " in values iterator.";
+
+ // Iterate over values for the current document ID
+ for (int i = 0, count = valuesIterator.valuesCount(); i < count; i++) {
+ long value = valuesIterator.nextValue();
+ // Assert that the value is as expected using the provided consumer
+ docCount += value;
+ }
+ }
+ }
+ return docCount;
+ }
+
+ public static XContentBuilder getExpandedMapping(int maxLeafDocs, boolean skipStarNodeCreationForSDVDimension) throws IOException {
+ return topMapping(b -> {
+ b.startObject("composite");
+ b.startObject("startree");
+ b.field("type", "star_tree");
+ b.startObject("config");
+ b.field("max_leaf_docs", maxLeafDocs);
+ if (skipStarNodeCreationForSDVDimension) {
+ b.startArray("skip_star_node_creation_for_dimensions");
+ b.value("sdv");
+ b.endArray();
+ }
+ b.startArray("ordered_dimensions");
+ b.startObject();
+ b.field("name", "sndv");
+ b.endObject();
+ b.startObject();
+ b.field("name", "sdv");
+ b.endObject();
+ b.startObject();
+ b.field("name", "dv");
+ b.endObject();
+ b.endArray();
+ b.startArray("metrics");
+ b.startObject();
+ b.field("name", "field");
+ b.startArray("stats");
+ b.value("sum");
+ b.value("value_count");
+ b.value("avg");
+ b.value("min");
+ b.value("max");
+ b.endArray();
+ b.endObject();
+ b.startObject();
+ b.field("name", "sndv");
+ b.startArray("stats");
+ b.value("sum");
+ b.value("value_count");
+ b.value("avg");
+ b.value("min");
+ b.value("max");
+ b.endArray();
+ b.endObject();
+ b.endArray();
+ b.endObject();
+ b.endObject();
+ b.endObject();
+ b.startObject("properties");
+ b.startObject("sndv");
+ b.field("type", "integer");
+ b.endObject();
+ b.startObject("sdv");
+ b.field("type", "integer");
+ b.endObject();
+ b.startObject("dv");
+ b.field("type", "integer");
+ b.endObject();
+ b.startObject("field");
+ b.field("type", "integer");
+ b.endObject();
+ b.endObject();
+ });
+ }
+}
From baecc1cfcc68c9e464a31dafffa46427e9d7d95c Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 21 Oct 2024 15:01:13 +0530
Subject: [PATCH 33/35] Rename valueCount() to entryValueCount()
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/common/util/FeatureFlags.java | 2 +-
.../startree/utils/StarTreeQueryHelper.java | 2 +-
.../SortedNumericStarTreeValuesIterator.java | 2 +-
.../aggregations/metrics/AvgAggregator.java | 2 +-
.../search/startree/StarTreeFilter.java | 4 +-
.../search/startree/StarTreeQueryContext.java | 1 -
.../search/SearchServiceStarTreeTests.java | 2 +
.../startree/StarTreeFilterTests.java | 2 +-
.../startree/StarTreeQueryContextTests.java | 127 ++++++++++++++++++
9 files changed, 136 insertions(+), 8 deletions(-)
create mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index 6df68013a8119..e663d8429da13 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
index b362b083a8d20..e538be5d5bece 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -212,7 +212,7 @@ public static LeafBucketCollector getStarTreeLeafCollector(
}
// Iterate over the values for the current entryId
- for (int i = 0, count = valuesIterator.valuesCount(); i < count; i++) {
+ for (int i = 0, count = valuesIterator.entryValueCount(); i < count; i++) {
long value = valuesIterator.nextValue();
valueConsumer.accept(value); // Apply the consumer operation (e.g., max, sum)
}
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
index 7d65e9eb1da8e..4b4bfa6a915eb 100644
--- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
+++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
@@ -30,7 +30,7 @@ public long nextValue() throws IOException {
return ((SortedNumericDocValues) docIdSetIterator).nextValue();
}
- public int valuesCount() throws IOException {
+ public int entryValueCount() throws IOException {
return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
}
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
index 0a3a514701e15..2970c5ca851e7 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
@@ -184,7 +184,7 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
}
// Iterate over the values for the current entryId
- for (int i = 0; i < sumValuesIterator.valuesCount(); i++) {
+ for (int i = 0; i < sumValuesIterator.entryValueCount(); i++) {
kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue()));
counts.increment(0, countValueIterator.nextValue()); // Apply the consumer operation (e.g., max, sum)
}
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
index f1cc1f11a671a..f7fa210691678 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
@@ -89,7 +89,7 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, Map<S
? bitSet.nextSetBit(entryId + 1)
: DocIdSetIterator.NO_MORE_DOCS) {
if (ndv.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) {
- final int valuesCount = ndv.valuesCount();
+ final int valuesCount = ndv.entryValueCount();
for (int i = 0; i < valuesCount; i++) {
long value = ndv.nextValue();
// Compare the value with the query value
@@ -207,7 +207,7 @@ private static StarTreeResult traverseStarTree(StarTreeValues starTreeValues, Ma
/**
* Helper class to wrap the result from traversing the star tree.
* */
- public static class StarTreeResult {
+ private static class StarTreeResult {
public final DocIdSetBuilder matchedDocIds;
public final Set<String> remainingPredicateColumns;
public final int numOfMatchedDocs;
diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
index 5a53dd122cdb0..cda3a25b30e53 100644
--- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
+++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java
@@ -76,5 +76,4 @@ public void setStarTreeValues(LeafReaderContext ctx, FixedBitSet values) {
starTreeValues[ctx.ord] = values;
}
}
-
}
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
index 0c88154ca2b38..f646b6b063ed0 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
@@ -141,6 +141,8 @@ private void assertStarTreeContext(
SearchContext context = searchService.createContext(reader, request, null, true);
StarTreeQueryContext actualContext = context.getStarTreeQueryContext();
+
+
if (expectedContext == null) {
assertThat(context.getStarTreeQueryContext(), nullValue());
} else {
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
index 009184d24817d..f8eb71a40319a 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java
@@ -242,7 +242,7 @@ private long getDocCountFromStarTree(CompositeIndexReader starTreeDocValuesReade
assert canAdvance : "Cannot advance to document ID " + bit + " in values iterator.";
// Iterate over values for the current document ID
- for (int i = 0, count = valuesIterator.valuesCount(); i < count; i++) {
+ for (int i = 0, count = valuesIterator.entryValueCount(); i < count; i++) {
long value = valuesIterator.nextValue();
// Assert that the value is as expected using the provided consumer
docCount += value;
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
new file mode 100644
index 0000000000000..076473d2c6a2a
--- /dev/null
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
@@ -0,0 +1,127 @@
+/*
+ * 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.aggregations.startree;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.util.FixedBitSet;
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.search.startree.StarTreeQueryContext;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class StarTreeQueryContextTests extends OpenSearchTestCase {
+
+ private CompositeIndexFieldInfo mockStarTree;
+ private LeafReaderContext mockLeafReaderContext;
+ private StarTreeQueryContext context;
+ private Map<String, Long> queryMap;
+
+ @Before
+ public void setUp() {
+ // Mock dependencies
+ mockStarTree = mock(CompositeIndexFieldInfo.class);
+ mockLeafReaderContext = mock(LeafReaderContext.class);
+
+ // Prepare queryMap
+ queryMap = new HashMap<>();
+ queryMap.put("field1", 123L);
+ queryMap.put("field2", 456L);
+
+ // Simulate leaf reader context ord
+ when(mockLeafReaderContext.ord).thenReturn(1);
+ }
+
+
+ public void testConstructorWithValidNumSegmentsCache() {
+ // Initialize context with valid numSegmentsCache
+ int numSegmentsCache = 3;
+ context = new StarTreeQueryContext(mockStarTree, queryMap, numSegmentsCache);
+
+ // Verify that the star tree and query map are set correctly
+ assertEquals(mockStarTree, context.getStarTree());
+ assertEquals(queryMap, context.getQueryMap());
+
+ // Verify that the starTreeValues array is initialized correctly
+ assertNotNull(context.getStarTreeValues());
+ assertEquals(numSegmentsCache, context.getStarTreeValues().length);
+ }
+
+
+ public void testConstructorWithNegativeNumSegmentsCache() {
+ // Initialize context with invalid numSegmentsCache (-1)
+ context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
+
+ // Verify that the starTreeValues array is not initialized
+ assertNull(context.getStarTreeValues());
+ }
+
+
+ public void testGetStarTreeValues_WithValidContext() {
+ // Initialize context with a cache of size 3
+ context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
+
+ // Create a FixedBitSet and assign it to the context
+ FixedBitSet fixedBitSet = new FixedBitSet(10);
+ context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
+
+ // Retrieve the FixedBitSet for the given context
+ FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
+
+ // Verify that the correct FixedBitSet is returned
+ assertNotNull(result);
+ assertEquals(fixedBitSet, result);
+ }
+
+
+ public void testGetStarTreeValues_WithNullCache() {
+ // Initialize context with no cache (numSegmentsCache is -1)
+ context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
+
+ // Retrieve the FixedBitSet for the given context
+ FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
+
+ // Verify that the result is null since there's no cache
+ assertNull(result);
+ }
+
+ public void testSetStarTreeValues() {
+ // Initialize context with a cache of size 3
+ context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
+
+ // Create a FixedBitSet and assign it to the context
+ FixedBitSet fixedBitSet = new FixedBitSet(10);
+ context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
+
+ // Retrieve the FixedBitSet for the given context and verify the update
+ FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
+ assertEquals(fixedBitSet, result);
+ }
+
+ public void testSetStarTreeValues_WithNullCache() {
+ // Initialize context with no cache (numSegmentsCache is -1)
+ context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
+
+ // Try setting a FixedBitSet and ensure no exception is thrown
+ FixedBitSet fixedBitSet = new FixedBitSet(10);
+ context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
+
+ // Since the cache is null, getStarTreeValues should return null
+ assertNull(context.getStarTreeValues(mockLeafReaderContext));
+ }
+}
From 949f95dfd3897ca01ec8ee98e9a3d6dbbddf3216 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 21 Oct 2024 15:34:09 +0530
Subject: [PATCH 34/35] spotless
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/search/SearchServiceStarTreeTests.java | 2 --
.../aggregations/startree/StarTreeQueryContextTests.java | 9 ++-------
2 files changed, 2 insertions(+), 9 deletions(-)
diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
index f646b6b063ed0..0c88154ca2b38 100644
--- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
+++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java
@@ -141,8 +141,6 @@ private void assertStarTreeContext(
SearchContext context = searchService.createContext(reader, request, null, true);
StarTreeQueryContext actualContext = context.getStarTreeQueryContext();
-
-
if (expectedContext == null) {
assertThat(context.getStarTreeQueryContext(), nullValue());
} else {
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
index 076473d2c6a2a..3d4b4719a09fa 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
+++ b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
@@ -10,11 +10,10 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.util.FixedBitSet;
-import org.junit.Before;
-import org.junit.Test;
import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
import org.opensearch.search.startree.StarTreeQueryContext;
import org.opensearch.test.OpenSearchTestCase;
+import org.junit.Before;
import java.util.HashMap;
import java.util.Map;
@@ -47,7 +46,6 @@ public void setUp() {
when(mockLeafReaderContext.ord).thenReturn(1);
}
-
public void testConstructorWithValidNumSegmentsCache() {
// Initialize context with valid numSegmentsCache
int numSegmentsCache = 3;
@@ -62,7 +60,6 @@ public void testConstructorWithValidNumSegmentsCache() {
assertEquals(numSegmentsCache, context.getStarTreeValues().length);
}
-
public void testConstructorWithNegativeNumSegmentsCache() {
// Initialize context with invalid numSegmentsCache (-1)
context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
@@ -71,7 +68,6 @@ public void testConstructorWithNegativeNumSegmentsCache() {
assertNull(context.getStarTreeValues());
}
-
public void testGetStarTreeValues_WithValidContext() {
// Initialize context with a cache of size 3
context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
@@ -88,7 +84,6 @@ public void testGetStarTreeValues_WithValidContext() {
assertEquals(fixedBitSet, result);
}
-
public void testGetStarTreeValues_WithNullCache() {
// Initialize context with no cache (numSegmentsCache is -1)
context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
@@ -99,7 +94,7 @@ public void testGetStarTreeValues_WithNullCache() {
// Verify that the result is null since there's no cache
assertNull(result);
}
-
+
public void testSetStarTreeValues() {
// Initialize context with a cache of size 3
context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
From 9742432150acd873ac26f28777152eef566b72b3 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 21 Oct 2024 20:34:11 +0530
Subject: [PATCH 35/35] delete tests added accidentally
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
.../opensearch/common/util/FeatureFlags.java | 2 +-
.../startree/StarTreeQueryContextTests.java | 122 ------------------
2 files changed, 1 insertion(+), 123 deletions(-)
delete mode 100644 server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
index e663d8429da13..6df68013a8119 100644
--- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
+++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -100,7 +100,7 @@ public class FeatureFlags {
* aggregations.
*/
public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled";
- public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);
+ public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
/**
* Gates the functionality of application based configuration templates.
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
deleted file mode 100644
index 3d4b4719a09fa..0000000000000
--- a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeQueryContextTests.java
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
- * 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.aggregations.startree;
-
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.util.FixedBitSet;
-import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
-import org.opensearch.search.startree.StarTreeQueryContext;
-import org.opensearch.test.OpenSearchTestCase;
-import org.junit.Before;
-
-import java.util.HashMap;
-import java.util.Map;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
-public class StarTreeQueryContextTests extends OpenSearchTestCase {
-
- private CompositeIndexFieldInfo mockStarTree;
- private LeafReaderContext mockLeafReaderContext;
- private StarTreeQueryContext context;
- private Map<String, Long> queryMap;
-
- @Before
- public void setUp() {
- // Mock dependencies
- mockStarTree = mock(CompositeIndexFieldInfo.class);
- mockLeafReaderContext = mock(LeafReaderContext.class);
-
- // Prepare queryMap
- queryMap = new HashMap<>();
- queryMap.put("field1", 123L);
- queryMap.put("field2", 456L);
-
- // Simulate leaf reader context ord
- when(mockLeafReaderContext.ord).thenReturn(1);
- }
-
- public void testConstructorWithValidNumSegmentsCache() {
- // Initialize context with valid numSegmentsCache
- int numSegmentsCache = 3;
- context = new StarTreeQueryContext(mockStarTree, queryMap, numSegmentsCache);
-
- // Verify that the star tree and query map are set correctly
- assertEquals(mockStarTree, context.getStarTree());
- assertEquals(queryMap, context.getQueryMap());
-
- // Verify that the starTreeValues array is initialized correctly
- assertNotNull(context.getStarTreeValues());
- assertEquals(numSegmentsCache, context.getStarTreeValues().length);
- }
-
- public void testConstructorWithNegativeNumSegmentsCache() {
- // Initialize context with invalid numSegmentsCache (-1)
- context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
-
- // Verify that the starTreeValues array is not initialized
- assertNull(context.getStarTreeValues());
- }
-
- public void testGetStarTreeValues_WithValidContext() {
- // Initialize context with a cache of size 3
- context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
-
- // Create a FixedBitSet and assign it to the context
- FixedBitSet fixedBitSet = new FixedBitSet(10);
- context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
-
- // Retrieve the FixedBitSet for the given context
- FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
-
- // Verify that the correct FixedBitSet is returned
- assertNotNull(result);
- assertEquals(fixedBitSet, result);
- }
-
- public void testGetStarTreeValues_WithNullCache() {
- // Initialize context with no cache (numSegmentsCache is -1)
- context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
-
- // Retrieve the FixedBitSet for the given context
- FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
-
- // Verify that the result is null since there's no cache
- assertNull(result);
- }
-
- public void testSetStarTreeValues() {
- // Initialize context with a cache of size 3
- context = new StarTreeQueryContext(mockStarTree, queryMap, 3);
-
- // Create a FixedBitSet and assign it to the context
- FixedBitSet fixedBitSet = new FixedBitSet(10);
- context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
-
- // Retrieve the FixedBitSet for the given context and verify the update
- FixedBitSet result = context.getStarTreeValues(mockLeafReaderContext);
- assertEquals(fixedBitSet, result);
- }
-
- public void testSetStarTreeValues_WithNullCache() {
- // Initialize context with no cache (numSegmentsCache is -1)
- context = new StarTreeQueryContext(mockStarTree, queryMap, -1);
-
- // Try setting a FixedBitSet and ensure no exception is thrown
- FixedBitSet fixedBitSet = new FixedBitSet(10);
- context.setStarTreeValues(mockLeafReaderContext, fixedBitSet);
-
- // Since the cache is null, getStarTreeValues should return null
- assertNull(context.getStarTreeValues(mockLeafReaderContext));
- }
-}