Integrated InstructionHandler and introduce delegation handle for data node execution

Signed-off-by: expani <anijainc@amazon.com>

Author: expani

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java

@@ -10,7 +10,9 @@
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.opensearch.analytics.spi.CommonExecutionContext;
+import org.opensearch.index.IndexSettings;
 import org.opensearch.index.engine.exec.IndexReaderProvider.Reader;
+import org.opensearch.index.mapper.MapperService;
 import org.opensearch.tasks.Task;
 
 /**
@@ -26,6 +28,8 @@ public class ShardScanExecutionContext implements CommonExecutionContext {
     private final Task task;
     private byte[] fragmentBytes;
     private BufferAllocator allocator;
+    private MapperService mapperService;
+    private IndexSettings indexSettings;
 
     /**
      * Constructs an execution context.
@@ -73,4 +77,24 @@ public BufferAllocator getAllocator() {
     public void setAllocator(BufferAllocator allocator) {
         this.allocator = allocator;
     }
+
+    /** Returns the shard's mapper service for field type resolution. */
+    public MapperService getMapperService() {
+        return mapperService;
+    }
+
+    /** Sets the shard's mapper service. */
+    public void setMapperService(MapperService mapperService) {
+        this.mapperService = mapperService;
+    }
+
+    /** Returns the shard's index settings. */
+    public IndexSettings getIndexSettings() {
+        return indexSettings;
+    }
+
+    /** Sets the shard's index settings. */
+    public void setIndexSettings(IndexSettings indexSettings) {
+        this.indexSettings = indexSettings;
+    }
 }

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

@@ -8,6 +8,8 @@
 
 package org.opensearch.analytics.spi;
 
+import java.util.List;
+
 /**
  * SPI extension point for backend query engine plugins.
  *
@@ -81,4 +83,34 @@ default ExchangeSinkProvider getExchangeSinkProvider() {
     default FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
         throw new UnsupportedOperationException("getInstructionHandlerFactory not implemented for [" + name() + "]");
     }
+
+    /**
+     * Prepare a filter delegation handle for the given delegated expressions.
+     * Called by Core after all instruction handlers have run, when the plan has delegation.
+     *
+     * <p>The accepting backend initializes its internal state (e.g., DirectoryReader,
+     * QueryShardContext, compiled Queries) and returns a handle that the driving backend
+     * will call into during execution.
+     *
+     * @param expressions the delegated expressions (annotationId + serialized query bytes)
+     * @param ctx the shared execution context (Reader, MapperService, IndexSettings)
+     * @return a handle the driving backend calls into via FFM upcalls
+     */
+    default FilterDelegationHandle getFilterDelegationHandle(List<DelegatedExpression> expressions, CommonExecutionContext ctx) {
+        throw new UnsupportedOperationException("getFilterDelegationHandle not implemented for [" + name() + "]");
+    }
+
+    /**
+     * Configure the driving backend to use the given delegation handle during execution.
+     * Called by Core after obtaining the handle from the accepting backend.
+     *
+     * <p>The driving backend registers the handle so that FFM upcalls from Rust
+     * (createProvider, createCollector, collectDocs) route to it.
+     *
+     * @param handle the delegation handle from the accepting backend
+     * @param backendContext the driving backend's execution context (from instruction handlers)
+     */
+    default void configureFilterDelegation(FilterDelegationHandle handle, BackendExecutionContext backendContext) {
+        throw new UnsupportedOperationException("configureFilterDelegation not implemented for [" + name() + "]");
+    }
 }

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

@@ -0,0 +1,58 @@
+/*
+ * 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.analytics.spi;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Describes the delegation metadata for a plan alternative. Carried on the wire
+ * alongside the instruction list so that Core can orchestrate the handle exchange
+ * between accepting and driving backends at the data node.
+ *
+ * @opensearch.internal
+ */
+public record DelegationDescriptor(
+    FilterTreeShape treeShape,
+    int delegatedPredicateCount,
+    List<DelegatedExpression> delegatedExpressions
+) implements Writeable {
+
+    public DelegationDescriptor(StreamInput in) throws IOException {
+        this(
+            in.readEnum(FilterTreeShape.class),
+            in.readVInt(),
+            readExpressions(in)
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeEnum(treeShape);
+        out.writeVInt(delegatedPredicateCount);
+        out.writeVInt(delegatedExpressions.size());
+        for (DelegatedExpression expr : delegatedExpressions) {
+            expr.writeTo(out);
+        }
+    }
+
+    private static List<DelegatedExpression> readExpressions(StreamInput in) throws IOException {
+        int count = in.readVInt();
+        List<DelegatedExpression> expressions = new ArrayList<>(count);
+        for (int i = 0; i < count; i++) {
+            expressions.add(new DelegatedExpression(in));
+        }
+        return expressions;
+    }
+}

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

@@ -0,0 +1,77 @@
+/*
+ * 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.analytics.spi;
+
+import java.io.Closeable;
+import java.lang.foreign.MemorySegment;
+
+/**
+ * Callback surface for filter delegation between a driving backend and an accepting backend.
+ *
+ * <p>One handle per query per shard. The accepting backend implements this interface;
+ * the driving backend calls into it via FFM upcalls during execution. Core closes it
+ * after execution completes.
+ *
+ * <p>Lifecycle:
+ * <ol>
+ *   <li>Rust calls {@link #createProvider(int)} once per delegated predicate (per annotationId)</li>
+ *   <li>Rust calls {@link #createCollector(int, int, int, int)} per (provider × segment)</li>
+ *   <li>Rust calls {@link #collectDocs(int, int, int, MemorySegment)} per row group</li>
+ *   <li>Rust calls {@link #releaseCollector(int)} when done with a segment</li>
+ *   <li>Rust calls {@link #releaseProvider(int)} when the query ends</li>
+ * </ol>
+ *
+ * @opensearch.internal
+ */
+public interface FilterDelegationHandle extends Closeable {
+
+    /**
+     * Create a provider for the given annotation ID. The accepting backend looks up
+     * the pre-compiled query for this annotation and prepares it for segment iteration.
+     *
+     * @param annotationId the annotation ID identifying the delegated predicate
+     * @return a provider key {@code >= 0}, or {@code -1} on failure
+     */
+    int createProvider(int annotationId);
+
+    /**
+     * Create a collector for one (segment, [minDoc, maxDoc)) range.
+     *
+     * @param providerKey key returned by {@link #createProvider(int)}
+     * @param segmentOrd the segment ordinal
+     * @param minDoc inclusive lower bound
+     * @param maxDoc exclusive upper bound
+     * @return a collector key {@code >= 0}, or {@code -1} on failure
+     */
+    int createCollector(int providerKey, int segmentOrd, int minDoc, int maxDoc);
+
+    /**
+     * Fill {@code out} with the matching doc-id bitset for the given collector.
+     *
+     * <p>Bit layout: word {@code i} contains matches for docs
+     * {@code [minDoc + i*64, minDoc + (i+1)*64)}, LSB-first within each word.
+     *
+     * @param collectorKey key returned by {@link #createCollector(int, int, int, int)}
+     * @param minDoc inclusive lower bound
+     * @param maxDoc exclusive upper bound
+     * @param out destination buffer; implementation writes up to {@code out.byteSize() / 8} words
+     * @return number of words written, or {@code -1} on error
+     */
+    int collectDocs(int collectorKey, int minDoc, int maxDoc, MemorySegment out);
+
+    /**
+     * Release resources for a collector.
+     */
+    void releaseCollector(int collectorKey);
+
+    /**
+     * Release resources for a provider.
+     */
+    void releaseProvider(int providerKey);
+}

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

@@ -44,7 +44,7 @@ public FilterDelegationInstructionNode(StreamInput in) throws IOException {
 
     @Override
     public InstructionType type() {
-        return InstructionType.SETUP_FILTER_DELEGATION_FOR_INDEX;
+        return InstructionType.SETUP_SHARD_SCAN_WITH_DELEGATION;
     }
 
     @Override

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

@@ -23,14 +23,14 @@ public enum InstructionType {
     /** Base scan setup — reader acquisition, SessionContext creation, default table provider. */
     SETUP_SHARD_SCAN,
     /**
-     * Filter delegation to an index backend — bridge setup, UDF registration, IndexedTableProvider.
+     * Filter delegation to an index backend — bridge setup, UDF registration, custom scan operator.
      *
      * <p>TODO: add a DelegationStrategy field (BACKEND_DRIVEN vs CENTRALLY_DRIVEN) to the
      * instruction node when centrally-driven delegation is implemented. Currently only
      * BACKEND_DRIVEN exists — derived from the backend declaring
      * {@code supportedDelegations(DelegationType.FILTER)}.
      */
-    SETUP_FILTER_DELEGATION_FOR_INDEX,
+    SETUP_SHARD_SCAN_WITH_DELEGATION,
     /** Partial aggregate mode — disable combine optimizer, cut plan to partial-only. */
     SETUP_PARTIAL_AGGREGATE,
     /** Final aggregate for coordinator reduce — ExchangeSink path, final-only agg. */
@@ -40,7 +40,7 @@ public enum InstructionType {
     public InstructionNode readNode(StreamInput in) throws IOException {
         return switch (this) {
             case SETUP_SHARD_SCAN -> new ShardScanInstructionNode(in);
-            case SETUP_FILTER_DELEGATION_FOR_INDEX -> new FilterDelegationInstructionNode(in);
+            case SETUP_SHARD_SCAN_WITH_DELEGATION -> new ShardScanWithDelegationInstructionNode(in);
             case SETUP_PARTIAL_AGGREGATE -> new PartialAggregateInstructionNode(in);
             case SETUP_FINAL_AGGREGATE -> new FinalAggregateInstructionNode(in);
         };

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

@@ -0,0 +1,59 @@
+/*
+ * 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.analytics.spi;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+/**
+ * Instruction node for shard scan with filter delegation — extends base shard scan
+ * with {@link FilterTreeShape} and delegated predicate count so the driving backend
+ * can configure its indexed execution path (UDF registration, IndexedTableProvider)
+ * in a single FFM call.
+ *
+ * @opensearch.internal
+ */
+public class ShardScanWithDelegationInstructionNode extends ShardScanInstructionNode {
+
+    private final FilterTreeShape treeShape;
+    private final int delegatedPredicateCount;
+
+    public ShardScanWithDelegationInstructionNode(FilterTreeShape treeShape, int delegatedPredicateCount) {
+        this.treeShape = treeShape;
+        this.delegatedPredicateCount = delegatedPredicateCount;
+    }
+
+    public ShardScanWithDelegationInstructionNode(StreamInput in) throws IOException {
+        super(in);
+        this.treeShape = in.readEnum(FilterTreeShape.class);
+        this.delegatedPredicateCount = in.readVInt();
+    }
+
+    @Override
+    public InstructionType type() {
+        return InstructionType.SETUP_SHARD_SCAN_WITH_DELEGATION;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeEnum(treeShape);
+        out.writeVInt(delegatedPredicateCount);
+    }
+
+    public FilterTreeShape getTreeShape() {
+        return treeShape;
+    }
+
+    public int getDelegatedPredicateCount() {
+        return delegatedPredicateCount;
+    }
+}

sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs

@@ -476,6 +476,27 @@ pub unsafe extern "C" fn df_create_session_context(
         .map_err(|e| e.to_string())
 }
 
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_create_session_context_indexed(
+    shard_view_ptr: i64,
+    runtime_ptr: i64,
+    table_name_ptr: *const u8,
+    table_name_len: i64,
+    context_id: i64,
+    tree_shape: i32,
+    delegated_predicate_count: i32,
+) -> i64 {
+    let table_name = str_from_raw(table_name_ptr, table_name_len)
+        .map_err(|e| format!("df_create_session_context_indexed: {}", e))?;
+    let mgr = get_rt_manager()?;
+    mgr.io_runtime
+        .block_on(crate::session_context::create_session_context_indexed(
+            runtime_ptr, shard_view_ptr, table_name, context_id, tree_shape, delegated_predicate_count,
+        ))
+        .map_err(|e| e.to_string())
+}
+
 #[ffm_safe]
 #[no_mangle]
 pub unsafe extern "C" fn df_cache_manager_remove_files(
@@ -618,11 +639,24 @@ pub unsafe extern "C" fn df_execute_with_context(
     let mgr = get_rt_manager()?;
     let plan_bytes = slice::from_raw_parts(plan_ptr, plan_len as usize);
     let cpu_executor = mgr.cpu_executor();
-    mgr.io_runtime
-        .block_on(crate::query_executor::execute_with_context(
-            session_handle,
-            plan_bytes,
-            cpu_executor,
-        ))
-        .map_err(|e| e.to_string())
+
+    // Route based on whether the session was configured for indexed execution
+    let handle_ref = &*(session_ctx_ptr as *const crate::session_context::SessionContextHandle);
+    if handle_ref.indexed_config.is_some() {
+        mgr.io_runtime
+            .block_on(crate::indexed_executor::execute_indexed_with_context(
+                session_ctx_ptr,
+                plan_bytes.to_vec(),
+                cpu_executor,
+            ))
+            .map_err(|e| e.to_string())
+    } else {
+        mgr.io_runtime
+            .block_on(crate::query_executor::execute_with_context(
+                session_handle,
+                plan_bytes,
+                cpu_executor,
+            ))
+            .map_err(|e| e.to_string())
+    }
 }