lookup and build engine bridge impl by engine name.

Signed-off-by: Marc Handalian <marc.handalian@gmail.com>

Author: mch2

plugins/engine-datafusion/jni/src/lib.rs

@@ -748,6 +748,8 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_streamNex
         with_jni_env(|env| {
             match result {
                 Ok(Some(batch)) => {
+                    log_info!("[RUST streamNext] Batch produced: {} rows, {} columns, schema: {:?}",
+                        batch.num_rows(), batch.num_columns(), batch.schema().fields().iter().map(|f| f.name().as_str()).collect::<Vec<_>>());
                     // Convert to FFI
                     let struct_array: StructArray = batch.into();
                     let array_data = struct_array.into_data();
@@ -756,6 +758,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_streamNex
                     set_action_listener_ok_global(env, &listener_ref, ffi_array_ptr as jlong);
                 }
                 Ok(None) => {
+                    log_info!("[RUST streamNext] End of stream reached");
                     // End of stream
                     set_action_listener_ok_global(env, &listener_ref, 0);
                 }

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

@@ -12,6 +12,7 @@
 import org.opensearch.analytics.backend.EngineBridge;
 import org.opensearch.analytics.spi.AnalyticsBackEndPlugin;
 import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
 
 /**
  * SPI adapter that delegates to the real {@link DataFusionPlugin} instance,
@@ -40,8 +41,8 @@ public String name() {
     }
 
     @Override
-    public EngineBridge<?, ?, ?> bridge(CatalogSnapshot snapshot) {
-        return delegate().bridge(snapshot);
+    public EngineBridge<?, ?, ?> bridge(CompositeEngine engine, CatalogSnapshot snapshot) {
+        return delegate().bridge(engine, snapshot);
     }
 
     @Override

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

@@ -31,6 +31,7 @@
 import org.opensearch.analytics.backend.EngineBridge;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.datafusion.jni.NativeBridge;
+import org.opensearch.datafusion.search.DatafusionContext;
 import org.opensearch.datafusion.search.DatafusionReader;
 
 import java.io.Closeable;

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

@@ -33,6 +33,7 @@
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.index.engine.SearchExecEngine;
 import org.opensearch.index.engine.exec.FileMetadata;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.ExtensiblePlugin;
@@ -212,7 +213,8 @@ public String name() {
     }
 
     @Override
-    public EngineBridge<?, ?, ?> bridge(CatalogSnapshot snapshot) {
+    public EngineBridge<?, ?, ?> bridge(CompositeEngine engine, CatalogSnapshot snapshot) {
+        DatafusionEngine dfEngine = (DatafusionEngine) engine.getEngine(name());
         long runtimePointer = dataFusionService.getRuntimePointer();
         Collection<WriterFileSet> files = snapshot.getSearchableFiles("parquet");
         // Derive directory path from the first WriterFileSet, or use empty string if no files

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsBackEndPlugin.java

@@ -11,6 +11,7 @@
 import org.apache.calcite.sql.SqlOperatorTable;
 import org.opensearch.analytics.backend.EngineBridge;
 import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
 
 
 /**
@@ -22,7 +23,7 @@ public interface AnalyticsBackEndPlugin {
     String name();
 
     /** JNI boundary for executing serialized plans, or null for engines without native execution. */
-    EngineBridge<?, ?, ?> bridge(CatalogSnapshot snapshot);
+    EngineBridge<?, ?, ?> bridge(CompositeEngine engine, CatalogSnapshot snapshot);
 
     /** Supported functions as a Calcite operator table, or null if the back-end adds no functions. */
     SqlOperatorTable operatorTable();

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackend.java

@@ -11,10 +11,12 @@
 import org.apache.calcite.sql.SqlOperatorTable;
 import org.opensearch.analytics.backend.EngineBridge;
 import org.opensearch.analytics.spi.AnalyticsBackEndPlugin;
+import org.opensearch.datafusion.DatafusionEngine;
 import org.opensearch.index.engine.exec.CatalogSnapshotAwareReaderManager;
 import org.opensearch.index.engine.exec.SearchExecEngine;
 import org.opensearch.index.engine.exec.WriterFileSet;
 import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.plugins.SearchAnalyticsBackEndPlugin;
 import org.opensearch.plugins.spi.vectorized.DataFormat;
@@ -45,14 +47,12 @@ public String name() {
     }
 
     @Override
-    public EngineBridge<?, ?, ?> bridge(CatalogSnapshot snapshot) {
-        Collection<WriterFileSet> files = snapshot.getSearchableFiles("parquet");
-        if (files.isEmpty() || files.stream().allMatch(wfs -> wfs.getFiles().isEmpty())) {
-            throw new IllegalStateException("No parquet files available in catalog snapshot");
+    public EngineBridge<?, ?, ?> bridge(CompositeEngine engine, CatalogSnapshot snapshot) {
+        try {
+            return new SandboxDataFusionBridge(getRuntimePtr(), (DatafusionReader) engine.getReader(name(), snapshot));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
         }
-        String dir = files.stream().findFirst().map(WriterFileSet::getDirectory).orElse("");
-        DatafusionReader reader = new DatafusionReader(dir, files);
-        return new SandboxDataFusionBridge(getRuntimePtr(), reader);
     }
 
     @Override

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchEnginePlugin.java

@@ -15,6 +15,8 @@
 import org.opensearch.index.engine.exec.CatalogSnapshotAwareReaderManager;
 import org.opensearch.index.engine.exec.IndexFilterProvider;
 import org.opensearch.index.engine.exec.SourceProvider;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CompositeEngine;
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.plugins.SearchAnalyticsBackEndPlugin;
 import org.opensearch.plugins.spi.vectorized.DataFormat;
@@ -36,7 +38,7 @@ public String name() {
     }
 
     @Override
-    public EngineBridge<?, ?, ?> bridge() {
+    public EngineBridge<?, ?, ?> bridge(CompositeEngine engine, CatalogSnapshot snapshot) {
         return null;
     }
 

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java

@@ -88,20 +88,21 @@ public Iterable<Object[]> execute(RelNode logicalFragment, Object context) {
 
         CompositeEngine engine = (CompositeEngine) indexShard.getIndexer();
 
-        // Prefer SearchExecEngine path if supported
-        if (plugin.supportsSearchExecEngine() && engine.getSearchBackendFactory() != null) {
-            logger.info("[DefaultPlanExecutor] Using SearchExecEngine path for back-end [{}]", plugin.name());
-            try {
-                return executeViaSearchExecEngine(engine, logicalFragment);
-            } catch (Exception e) {
-                throw new RuntimeException("SearchExecEngine execution failed for [" + plugin.name() + "]", e);
-            }
-        }
+//        // Prefer SearchExecEngine path if supported
+//        if (plugin.supportsSearchExecEngine() && engine.getSearchBackendFactory() != null) {
+//            logger.info("[DefaultPlanExecutor] Using SearchExecEngine path for back-end [{}]", plugin.name());
+//            try {
+//                return executeViaSearchExecEngine(engine, logicalFragment);
+//            } catch (Exception e) {
+//                throw new RuntimeException("SearchExecEngine execution failed for [" + plugin.name() + "]", e);
+//            }
+//        }
 
         // Bridge path
         try (CompositeEngine.ReleasableRef<CatalogSnapshot> snapshot = engine.acquireSnapshot()) {
             EngineBridge<byte[], ? extends EngineResultStream, RelNode> bridge =
-                (EngineBridge<byte[], ? extends EngineResultStream, RelNode>) plugin.bridge(snapshot.getRef());
+                (EngineBridge<byte[], ? extends EngineResultStream, RelNode>) plugin.bridge(engine, snapshot.getRef());
+
             byte[] converted = bridge.convertFragment(logicalFragment);
 
             List<Object[]> rows = new ArrayList<>();

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ClickBenchRestIT.java

@@ -117,6 +117,17 @@ private void bulkInsertTestData() throws IOException {
 //        logger.info("Bulk inserted test data into {} index", HITS_INDEX);
     }
 
+    // Expected results for each query based on the 10-doc test dataset.
+    // Each entry: queryId -> expected single-row values (in column order).
+    private static final Map<String, List<Object>> EXPECTED_RESULTS = Map.of(
+        "q1", List.of(10),           // COUNT(*)
+        "q2", List.of(2),            // COUNT(*) WHERE AdvEngineID != 0  (docs 1,7)
+        "q3", List.of(5, 10),        // SUM(AdvEngineID)=5, COUNT(*)=10
+        "q4", List.of(435090932899640449L), // MAX(UserID)
+        "q5", List.of(10),           // COUNT(DISTINCT UserID) - all 10 unique
+        "q6", List.of(7)             // COUNT(DISTINCT SearchPhrase) - 6 non-empty + 1 empty = 7
+    );
+
     public void testClickBenchQueries() throws Exception {
         createHitsIndex();
         bulkInsertTestData();
@@ -164,6 +175,29 @@ private void runClickBenchQuery(String queryId) throws Exception {
         List<String> columns = (List<String>) responseMap.get("columns");
         assertFalse("Columns should not be empty for " + queryId, columns.isEmpty());
 
+        @SuppressWarnings("unchecked")
+        List<List<Object>> rows = (List<List<Object>>) responseMap.get("rows");
+        logger.info("RESULT {}: columns={}, rows={}", queryId, columns, rows.size());
+        for (int i = 0; i < Math.min(rows.size(), 10); i++) {
+            logger.info("  row[{}]: {}", i, rows.get(i));
+        }
+        if (rows.size() > 10) {
+            logger.info("  ... ({} more rows)", rows.size() - 10);
+        }
+
+        // Assert expected results
+        List<Object> expected = EXPECTED_RESULTS.get(queryId);
+        if (expected != null) {
+            assertEquals("Expected exactly 1 row for " + queryId, 1, rows.size());
+            List<Object> actualRow = rows.get(0);
+            assertEquals("Column count mismatch for " + queryId, expected.size(), actualRow.size());
+            for (int i = 0; i < expected.size(); i++) {
+                long expectedVal = ((Number) expected.get(i)).longValue();
+                long actualVal = ((Number) actualRow.get(i)).longValue();
+                assertEquals(queryId + " column " + columns.get(i) + " (index " + i + ")", expectedVal, actualVal);
+            }
+        }
+
         logger.info("SUCCESS {}: {} columns", queryId, columns.size());
     }
 

server/src/main/java/org/opensearch/index/engine/exec/coord/CompositeEngine.java

@@ -97,6 +97,7 @@
 import org.opensearch.index.translog.listener.TranslogEventListener;
 import org.opensearch.indices.pollingingest.PollingIngestStats;
 import org.opensearch.plugins.PluginsService;
+import org.opensearch.plugins.SearchAnalyticsBackEndPlugin;
 import org.opensearch.plugins.SearchEnginePlugin;
 import org.opensearch.search.suggest.completion.CompletionStats;
 import org.opensearch.plugins.spi.vectorized.DataFormat;
@@ -173,6 +174,8 @@ public class CompositeEngine implements LifecycleAware, Closeable, Indexer, Chec
     private final List<CatalogSnapshotAwareRefreshListener> catalogSnapshotAwareRefreshListeners = new ArrayList<>();
     private final Map<String, List<FileDeletionListener>> fileDeletionListeners = new HashMap<>();
     private final Map<DataFormat, List<SearchExecEngine<?, ?, ?, ?>>> readEngines = new HashMap<>();
+    private final Map<String, SearchExecEngine<?, ?, ?, ?>> readEnginesByName = new HashMap<>();
+    private final Map<String, CatalogSnapshotAwareReaderManager<?>> readerManagersByName = new HashMap<>();
     @Nullable
     private final SearchBackendFactory searchBackendFactory;
     private final MergeScheduler mergeScheduler;
@@ -406,6 +409,7 @@ public void onFailure(String reason, Exception ex) {
 
                     currentSearchEngines.add(newSearchEngine);
                     readEngines.put(dataFormat, currentSearchEngines);
+                    readEnginesByName.put(dataFormat.getName(), newSearchEngine);
 
                     // TODO : figure out how to do internal and external refresh listeners
                     // Maybe external refresh should be managed in opensearch core and plugins should always give
@@ -434,6 +438,17 @@ public void onFailure(String reason, Exception ex) {
             }
             this.searchBackendFactory = backendFactory;
 
+            // Populate name-keyed reader managers from SearchAnalyticsBackEndPlugins
+            for (SearchAnalyticsBackEndPlugin backendPlugin : pluginsService.filterPlugins(SearchAnalyticsBackEndPlugin.class)) {
+                for (DataFormat format : backendPlugin.getSupportedFormats()) {
+                    try {
+                        readerManagersByName.put(backendPlugin.name(), backendPlugin.createReaderManager(format, shardPath));
+                    } catch (Exception e) {
+                        logger.warn("Failed to create reader manager for plugin [{}], format [{}]", backendPlugin.name(), format, e);
+                    }
+                }
+            }
+
             success = true;
         } catch (IOException | TranslogCorruptedException e) {
             throw new EngineCreationFailureException(shardId, "failed to create engine", e);
@@ -605,6 +620,31 @@ public Object getReader(DataFormat format, CatalogSnapshot snapshot) throws IOEx
         return rm != null ? rm.getReader(snapshot) : null;
     }
 
+    public Object getEngine(String name) {
+        return readEnginesByName.get(name);
+    }
+
+    /**
+     * Returns the catalog-snapshot-aware reader manager for the given name, or null.
+     */
+    @Nullable
+    public CatalogSnapshotAwareReaderManager<?> getReaderManager(String name) {
+        return readerManagersByName.get(name);
+    }
+
+    /**
+     * Returns a reader for the given name and catalog snapshot, or null.
+     */
+    @Nullable
+    public Object getReader(String name, CatalogSnapshot snapshot) throws IOException {
+        CatalogSnapshotAwareReaderManager<?> rm = getReaderManager(name);
+        // hack
+        if (rm != null) {
+            rm.afterRefresh(true, snapshot);
+        }
+        return rm != null ? rm.getReader(snapshot) : null;
+    }
+
     /**
      * Returns a reader for the given format using the current catalog snapshot.
      * TODO: return composite reader that spans all formats