Add changes for searcher integration

Co-authored-by: Arpit Bandejiya <[email protected]>
Signed-off-by: Arpit Bandejiya <[email protected]>

Author: bharath-techie

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionPlugin.java

@@ -125,7 +125,7 @@ public List<DataFormat> getSupportedFormats() {
     public SearchExecEngine<DatafusionContext, DatafusionSearcher,
             DatafusionReaderManager, DatafusionQuery>
         createEngine(DataFormat dataFormat,Collection<FileMetadata> formatCatalogSnapshot) throws IOException {
-        return new DatafusionEngine(dataFormat, formatCatalogSnapshot);
+        return new DatafusionEngine(dataFormat, formatCatalogSnapshot, dataFusionService);
     }
 
     /**

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java

@@ -161,6 +161,9 @@ public CompletableFuture<RecordBatchStream> executeSubstraitQuery(long sessionCo
         return engine.executeSubstraitQuery(sessionContextId, substraitPlanBytes);
     }
 
+    public long getRuntimePointer() {
+        return globalRuntimeEnv.getPointer();
+    }
     /**
      * Close the session context and clean up resources
      *

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java

@@ -8,8 +8,15 @@
 
 package org.opensearch.datafusion;
 
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.action.search.SearchShardTask;
 import org.opensearch.common.lease.Releasables;
+import org.opensearch.datafusion.core.DefaultRecordBatchStream;
 import org.opensearch.datafusion.search.DatafusionContext;
 import org.opensearch.datafusion.search.DatafusionQuery;
 import org.opensearch.datafusion.search.DatafusionQueryPhaseExecutor;
@@ -23,6 +30,7 @@
 import org.opensearch.index.engine.EngineSearcherSupplier;
 import org.opensearch.index.engine.SearchExecEngine;
 import org.opensearch.index.engine.exec.FileMetadata;
+import org.opensearch.search.aggregations.SearchResultsCollector;
 import org.opensearch.search.internal.ReaderContext;
 import org.opensearch.search.internal.ShardSearchRequest;
 import org.opensearch.search.query.QueryPhaseExecutor;
@@ -33,17 +41,23 @@
 import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.function.Function;
 
 public class DatafusionEngine extends SearchExecEngine<DatafusionContext, DatafusionSearcher,
     DatafusionReaderManager, DatafusionQuery> {
 
+    private static final Logger logger = LogManager.getLogger(DatafusionEngine.class);
+
     private DataFormat dataFormat;
     private DatafusionReaderManager datafusionReaderManager;
+    private DataFusionService datafusionService;
 
-    public DatafusionEngine(DataFormat dataFormat, Collection<FileMetadata> formatCatalogSnapshot) throws IOException {
+    public DatafusionEngine(DataFormat dataFormat, Collection<FileMetadata> formatCatalogSnapshot, DataFusionService dataFusionService) throws IOException {
         this.dataFormat = dataFormat;
         this.datafusionReaderManager = new DatafusionReaderManager("TODO://FigureOutPath", formatCatalogSnapshot);
+        this.datafusionService = dataFusionService;
     }
 
     @Override
@@ -60,7 +74,7 @@ public QueryPhaseExecutor<DatafusionContext> getQueryPhaseExecutor() {
     public DatafusionContext createContext(ReaderContext readerContext, ShardSearchRequest request, SearchShardTask task) throws IOException {
         DatafusionContext datafusionContext = new DatafusionContext(readerContext, request, task, this);
         // Parse source
-        datafusionContext.datafusionQuery(new DatafusionQuery(request.source().getSubstraitBytes(), new ArrayList<>()));
+        datafusionContext.datafusionQuery(new DatafusionQuery(request.source().queryPlanIR(), new ArrayList<>()));
         return datafusionContext;
     }
 
@@ -138,4 +152,48 @@ public CatalogSnapshotAwareRefreshListener getRefreshListener(Engine.SearcherSco
     public boolean assertSearcherIsWarmedUp(String source, Engine.SearcherScope scope) {
         return false;
     }
+
+    @Override
+    public Map<String, Object[]> execute(DatafusionContext context) {
+
+        Map<String, Object[]> finalRes = new HashMap<>();
+        try {
+            DatafusionSearcher datafusionSearcher = context.getEngineSearcher();
+            long streamPointer = datafusionSearcher.search(context.getDatafusionQuery());
+            RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+            RecordBatchStream stream = new RecordBatchStream(streamPointer, datafusionService.getRuntimePointer() , allocator);
+
+            // We can have some collectors passed like this which can collect the results and convert to InternalAggregation
+            // Is the possible? need to check
+
+            SearchResultsCollector<RecordBatchStream> collector = new SearchResultsCollector<RecordBatchStream>() {
+                @Override
+                public void collect(RecordBatchStream value) {
+                    VectorSchemaRoot root = value.getVectorSchemaRoot();
+                    for (Field field : root.getSchema().getFields()) {
+                        String filedName = field.getName();
+                        FieldVector fieldVector = root.getVector(filedName);
+                        Object[] fieldValues = new Object[fieldVector.getValueCount()];
+                        for (int i = 0; i < fieldVector.getValueCount(); i++) {
+                            fieldValues[i] = fieldVector.getObject(i);
+                        }
+                        finalRes.put(filedName, fieldValues);
+                    }
+                }
+            };
+
+            while (stream.loadNextBatch().join()) {
+                collector.collect(stream);
+            }
+
+            logger.info("Final Results:");
+            for (Map.Entry<String, Object[]> entry : finalRes.entrySet()) {
+                logger.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue()));
+            }
+
+        } catch (Exception exception) {
+            logger.error("Failed to execute Substrait query plan", exception);
+        }
+        return finalRes;
+    }
 }

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/RecordBatchStream.java

@@ -29,23 +29,22 @@
  */
 public class RecordBatchStream {
 
-    private final SessionContext context;
     private final long streamPointer;
     private final BufferAllocator allocator;
     private final CDataDictionaryProvider dictionaryProvider;
     private boolean initialized = false;
     private VectorSchemaRoot vectorSchemaRoot = null;
+    private long runtimePtr;
 
     /**
      * Creates a new RecordBatchStream for the given stream pointer
-     * @param ctx the session context
-     * @param streamId pointer to the native stream
+     * @param streamId the stream pointer
      * @param allocator memory allocator for Arrow vectors
      */
-    public RecordBatchStream(SessionContext ctx, long streamId, BufferAllocator allocator) {
-        this.context = ctx;
+    public RecordBatchStream(long streamId, long runtimePtr, BufferAllocator allocator) {
         this.streamPointer = streamId;
         this.allocator = allocator;
+        this.runtimePtr = runtimePtr;
         this.dictionaryProvider = new CDataDictionaryProvider();
     }
 
@@ -99,7 +98,7 @@ private void ensureInitialized() {
      */
     public CompletableFuture<Boolean> loadNextBatch() {
         ensureInitialized();
-        long runtimePointer = context.getRuntime();
+        long runtimePointer = this.runtimePtr;
         CompletableFuture<Boolean> result = new CompletableFuture<>();
         next(runtimePointer, streamPointer, (errString, arrowArrayAddress) -> {
             if (ErrorUtil.containsError(errString)) {

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

@@ -48,6 +48,11 @@ public void search(DatafusionQuery datafusionQuery, List<SearchResultsCollector<
         }
     }
 
+    @Override
+    public long search(DatafusionQuery datafusionQuery) {
+        return DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.getSubstraitBytes());
+    }
+
     public DatafusionReader getReader() {
         return reader;
     }

server/src/main/java/org/opensearch/index/engine/EngineSearcher.java

@@ -13,6 +13,7 @@
 import org.opensearch.search.aggregations.SearchResultsCollector;
 
 import java.io.IOException;
+import java.io.UnsupportedEncodingException;
 import java.util.List;
 
 @ExperimentalApi
@@ -30,4 +31,8 @@ public interface EngineSearcher<Q,C> extends Releasable {
     default void search(Q query, List<SearchResultsCollector<C>> collectors) throws IOException {
         throw new UnsupportedOperationException();
     }
+
+    default long search(Q query) throws IOException {
+        throw new UnsupportedOperationException();
+    }
 }

server/src/main/java/org/opensearch/index/engine/SearchExecEngine.java

@@ -17,6 +17,7 @@
 import org.opensearch.search.query.QueryPhaseExecutor;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  * Generic read engine interface that provides searcher operations and query phase execution
@@ -43,4 +44,10 @@ public abstract class SearchExecEngine<C extends SearchContext, S extends Engine
      * Create a search context for this engine
      */
     public abstract C createContext(ReaderContext readerContext, ShardSearchRequest request, SearchShardTask task) throws IOException;
+
+    /**
+     * execute
+     * @return
+     */
+    public abstract Map<String, Object[]> execute(C context) throws IOException;
 }

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

@@ -813,7 +813,8 @@ private SearchPhaseResult executeQueryPhase(
         // Till here things are generic but for datafusion , we need to abstract out and get the read engine specific implementation
         // it could be reusing existing
         final ReaderContext readerContext = createOrGetReaderContext(request, keepStatesInContext);
-        SearchExecEngine<?, ?, ?, ?> searchExecEngine = readerContext.indexShard()
+        @SuppressWarnings("unchecked")
+        SearchExecEngine searchExecEngine = readerContext.indexShard()
             .getIndexingExecutionCoordinator()
             .getPrimaryReadEngine();
 
@@ -828,10 +829,10 @@ private SearchPhaseResult executeQueryPhase(
         ) {
 
             // TODO Execute plan here
+            // TODO : figure out how to tie this
             byte[] substraitQuery = request.source().queryPlanIR();
             if (substraitQuery != null) {
-                SearchExecutionEngine searchExecutionEngine = readerContext.indexShard().getSearchExecutionEngine();
-                Map<String, Object[]> result = searchExecutionEngine.execute(substraitQuery);
+                Map<String, Object[]> result = searchExecEngine.execute(context);
                 context.setDFResults(result);
             }
 
@@ -841,9 +842,14 @@ private SearchPhaseResult executeQueryPhase(
             }
             final long afterQueryTime;
             try (SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(context)) {
+                // TODO check for this
+//                @SuppressWarnings("unchecked")
+//                QueryPhaseExecutor<SearchContext> queryPhaseExecutor =
+//                    (QueryPhaseExecutor<SearchContext>) searchExecEngine.getQueryPhaseExecutor();
+
                 //QueryPhaseExecutor<?> queryPhaseExecutor = readEngine.getQueryPhaseExecutor();
-                //boolean success = queryPhaseExecutor.execute(context);
-                //loadOrExecuteQueryPhase(request, context);
+//                boolean success = queryPhaseExecutor.execute(context);
+                loadOrExecuteQueryPhase(request, context);
                 queryPhase.execute(context);
                 // loadOrExecuteQueryPhase(request, context);
                 if (context.queryResult().hasSearchContext() == false && readerContext.singleSession()) {