Changes to make plugin contexts work with source parse

Signed-off-by: bharath-techie <[email protected]>

Author: bharath-techie

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionContext.java

@@ -76,6 +76,8 @@ public class DatafusionContext extends SearchContext {
     private final QueryShardContext queryShardContext;
     private DatafusionQuery datafusionQuery;
     private Map<String, Object[]> dfResults;
+    private SearchContextAggregations aggregations;
+
     /**
      * Constructor
      * @param readerContext The reader context
@@ -224,7 +226,7 @@ public ScrollContext scrollContext() {
 
     @Override
     public SearchContextAggregations aggregations() {
-        return null;
+        return aggregations;
     }
 
     /**
@@ -233,7 +235,8 @@ public SearchContextAggregations aggregations() {
      */
     @Override
     public SearchContext aggregations(SearchContextAggregations aggregations) {
-        return null;
+        this.aggregations = aggregations;
+        return this;
     }
 
     /**
@@ -790,4 +793,5 @@ public void setDFResults(Map<String, Object[]> dfResults) {
     public Map<String, Object[]> getDFResults() {
         return dfResults;
     }
+
 }

server/src/main/java/org/opensearch/search/SearchService.java

@@ -825,15 +825,15 @@ private SearchPhaseResult executeQueryPhase(
         );
         try (
             Releasable ignored = readerContext.markAsUsed(getKeepAlive(request));
-            //SearchContext context = createContext(readerContext, request, task, true, isStreamSearch)
-
             // Get engine-specific executor and context
             // TODO : move this logic to work with Lucene
 
-            SearchContext context = searchExecEngine.createContext(readerContext, request, shardTarget, task);
+            SearchContext context = createContext(readerContext, request, task, true, isStreamSearch, searchExecEngine);
+
             //SearchContext context = createContext(readerContext, request, task, true)
         ) {
-
+            // TODO : this is not correct - need to tie source to plugin context above
+            //context.aggregations(context1.aggregations());
             // TODO Execute plan here
             // TODO : figure out how to tie this
             byte[] substraitQuery = request.source().queryPlanIR();
@@ -1270,17 +1270,26 @@ final SearchContext createContext(
         SearchShardTask task,
         boolean includeAggregations
     ) throws IOException {
-        return createContext(readerContext, request, task, includeAggregations, false);
+        return createContext(readerContext, request, task, includeAggregations, false, null);
     }
 
     private SearchContext createContext(
         ReaderContext readerContext,
         ShardSearchRequest request,
         SearchShardTask task,
         boolean includeAggregations,
-        boolean isStreamSearch
+        boolean isStreamSearch,
+        SearchExecEngine searchExecEngine
     ) throws IOException {
-        final DefaultSearchContext context = createSearchContext(readerContext, request, defaultSearchTimeout, false, isStreamSearch);
+        //final DefaultSearchContext context = createSearchContext(readerContext, request, defaultSearchTimeout, false, isStreamSearch);
+
+        SearchShardTarget shardTarget = new SearchShardTarget(
+            clusterService.localNode().getId(),
+            readerContext.indexShard().shardId(),
+            request.getClusterAlias(),
+            OriginalIndices.NONE
+        );
+        SearchContext context = searchExecEngine.createContext(readerContext, request, shardTarget, task);
         try {
             if (request.scroll() != null) {
                 context.scrollContext().scroll = request.scroll();
@@ -1546,10 +1555,10 @@ private void processFailure(ReaderContext context, Exception exc) {
         }
     }
 
-    private void parseSource(DefaultSearchContext context, SearchSourceBuilder source, boolean includeAggregations) {
+    private void parseSource(SearchContext context, SearchSourceBuilder source, boolean includeAggregations) {
         // nothing to parse...
         if (source == null) {
-            context.evaluateRequestShouldUseConcurrentSearch();
+           // context.evaluateRequestShouldUseConcurrentSearch(); // TODO : specific to default search context
             return;
         }
 
@@ -1706,7 +1715,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             if (context.scrollContext() == null && !(context.readerContext() instanceof PitReaderContext)) {
                 throw new SearchException(shardTarget, "`slice` cannot be used outside of a scroll context or PIT context");
             }
-            context.sliceBuilder(source.slice());
+            // context.sliceBuilder(source.slice());  // TODO : specific to default search context
         }
 
         if (source.storedFields() != null) {
@@ -1740,13 +1749,13 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             final CollapseContext collapseContext = source.collapse().build(queryShardContext);
             context.collapse(collapseContext);
         }
-        context.evaluateRequestShouldUseConcurrentSearch();
+        // context.evaluateRequestShouldUseConcurrentSearch();  // TODO : specific to default search context
         if (source.profile()) {
             final Function<Query, Collection<Supplier<ProfileMetric>>> pluginProfileMetricsSupplier = (query) -> pluginProfilers.stream()
                 .flatMap(p -> p.getQueryProfileMetrics(context, query).stream())
                 .toList();
             Profilers profilers = new Profilers(context.searcher(), context.shouldUseConcurrentSearch(), pluginProfileMetricsSupplier);
-            context.setProfilers(profilers);
+            // context.setProfilers(profilers); // TODO : specific to default search context
         }
 
         if (context.getStarTreeIndexEnabled() && StarTreeQueryHelper.isStarTreeSupported(context)) {