use rewrite instead of createweight in originalorstartreequery

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

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;
     }
 

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;
         }

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;
     }
 }

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;

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

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,15 +56,19 @@
 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.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;
@@ -72,6 +78,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;
@@ -110,6 +119,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;
@@ -2082,4 +2092,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));
+            }
+        }
+    }
+
 }

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,