getDelegatedQueries() {
+ return delegatedQueries;
+ }
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterTreeShape.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterTreeShape.java
new file mode 100644
index 0000000000000..8081ba7d63cb6
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterTreeShape.java
@@ -0,0 +1,32 @@
+/*
+ * 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;
+
+/**
+ * Backend-agnostic description of the boolean tree shape when filter delegation is active.
+ * Provided by the planner so backends can choose their execution strategy without
+ * re-inspecting the Substrait plan.
+ *
+ * @opensearch.internal
+ */
+public enum FilterTreeShape {
+ /** No delegation — all predicates handled natively by the driving backend. */
+ NO_DELEGATION,
+ /**
+ * All predicates (delegated + native) are under a single AND — no interleaving
+ * under OR/NOT. Backend can handle delegated bitsets and native predicates independently.
+ */
+ CONJUNCTIVE,
+ /**
+ * Delegated and native predicates are interleaved under OR/NOT — the boolean tree
+ * mixes predicates from different backends under non-AND operators. Backend needs a
+ * tree evaluator to combine bitsets from both backends per the boolean structure.
+ */
+ INTERLEAVED_BOOLEAN_EXPRESSION
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FinalAggregateInstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FinalAggregateInstructionNode.java
new file mode 100644
index 0000000000000..87bfc2c5081d8
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FinalAggregateInstructionNode.java
@@ -0,0 +1,41 @@
+/*
+ * 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 final aggregate in coordinator reduce — ExchangeSink path,
+ * remove partial agg, preserve final-only for the driving backend's reduce execution.
+ *
+ * TODO: add backend-specific config fields as final aggregate implementation is built out.
+ *
+ * @opensearch.internal
+ */
+public class FinalAggregateInstructionNode implements InstructionNode {
+
+ public FinalAggregateInstructionNode() {}
+
+ public FinalAggregateInstructionNode(StreamInput in) throws IOException {
+ // TODO: read config fields when added
+ }
+
+ @Override
+ public InstructionType type() {
+ return InstructionType.SETUP_FINAL_AGGREGATE;
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ // TODO: write config fields when added
+ }
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandler.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandler.java
new file mode 100644
index 0000000000000..db70c1c9fdd33
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandler.java
@@ -0,0 +1,31 @@
+/*
+ * 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;
+
+/**
+ * Applies an {@link InstructionNode} to the execution context at the data node.
+ * Each handler is created per-execution by the backend's
+ * {@link FragmentInstructionHandlerFactory#createHandler(InstructionNode)}.
+ *
+ * @param the concrete instruction node type this handler processes
+ * @opensearch.internal
+ */
+public interface FragmentInstructionHandler {
+
+ /**
+ * Applies the instruction, reading from Core's context and building upon the
+ * backend's accumulated execution context from previous handlers.
+ *
+ * @param node the instruction node
+ * @param commonContext Core-provided context (shard info or reduce info)
+ * @param backendContext backend state from previous handler, or {@code null} for the first handler
+ * @return updated backend execution context for the next handler or final consumer
+ */
+ BackendExecutionContext apply(N node, CommonExecutionContext commonContext, BackendExecutionContext backendContext);
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java
new file mode 100644
index 0000000000000..6e62fa10f012e
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java
@@ -0,0 +1,52 @@
+/*
+ * 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.util.List;
+import java.util.Optional;
+
+/**
+ * Factory for creating {@link InstructionNode}s at the coordinator and
+ * {@link FragmentInstructionHandler}s at the data node. One factory per backend,
+ * accessed via {@code AnalyticsSearchBackendPlugin.getInstructionHandlerFactory()}.
+ *
+ * Coordinator-side creation methods return {@link Optional#empty()} if the backend
+ * does not support the instruction type. Core logs and skips unsupported instructions.
+ *
+ * @opensearch.internal
+ */
+public interface FragmentInstructionHandlerFactory {
+
+ // ── Coordinator-side: create instruction nodes ──
+
+ /** Creates a shard scan instruction node. */
+ Optional createShardScanNode();
+
+ /** Creates a filter delegation instruction node with the given delegation metadata. */
+ Optional createFilterDelegationNode(
+ FilterTreeShape treeShape,
+ int delegatedPredicateCount,
+ List delegatedQueries
+ );
+
+ /** Creates a partial aggregate instruction node. */
+ Optional createPartialAggregateNode();
+
+ /** Creates a final aggregate instruction node for coordinator reduce. */
+ Optional createFinalAggregateNode();
+
+ // ── Data-node-side: create handler for an instruction node ──
+
+ /**
+ * Creates a handler for the given instruction node. The handler's
+ * {@link FragmentInstructionHandler#apply} will be called with the node
+ * and the execution context.
+ */
+ FragmentInstructionHandler createHandler(InstructionNode node);
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionNode.java
new file mode 100644
index 0000000000000..e52e545d0384b
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionNode.java
@@ -0,0 +1,27 @@
+/*
+ * 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.Writeable;
+
+/**
+ * Metadata node produced by the planner (via backend's factory) at the coordinator
+ * and consumed by the backend's handler at the data node. Carries typed configuration
+ * that the handler uses to configure the execution environment.
+ *
+ * Generic parent interface — backends extend with concrete classes if they need
+ * additional coordinator-side context beyond what the framework provides.
+ *
+ * @opensearch.internal
+ */
+public interface InstructionNode extends Writeable {
+
+ /** The instruction type — used to look up the handler factory at the data node. */
+ InstructionType type();
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionType.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionType.java
new file mode 100644
index 0000000000000..490f60a967707
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/InstructionType.java
@@ -0,0 +1,48 @@
+/*
+ * 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 java.io.IOException;
+
+/**
+ * Types of instructions that the planner can produce for backend execution.
+ * Each type corresponds to a specific execution concern that the backend
+ * must handle during the prepare phase on the data node.
+ *
+ * @opensearch.internal
+ */
+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.
+ *
+ *
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,
+ /** Partial aggregate mode — disable combine optimizer, cut plan to partial-only. */
+ SETUP_PARTIAL_AGGREGATE,
+ /** Final aggregate for coordinator reduce — ExchangeSink path, final-only agg. */
+ SETUP_FINAL_AGGREGATE;
+
+ /** Deserializes an {@link InstructionNode} from the stream based on this type. */
+ 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_PARTIAL_AGGREGATE -> new PartialAggregateInstructionNode(in);
+ case SETUP_FINAL_AGGREGATE -> new FinalAggregateInstructionNode(in);
+ };
+ }
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/PartialAggregateInstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/PartialAggregateInstructionNode.java
new file mode 100644
index 0000000000000..2f94d08f3ef0f
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/PartialAggregateInstructionNode.java
@@ -0,0 +1,40 @@
+/*
+ * 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 partial aggregate mode — disable combine optimizer, cut plan to partial-only.
+ *
+ *
TODO: add backend-specific config fields as partial aggregate implementation is built out.
+ *
+ * @opensearch.internal
+ */
+public class PartialAggregateInstructionNode implements InstructionNode {
+
+ public PartialAggregateInstructionNode() {}
+
+ public PartialAggregateInstructionNode(StreamInput in) throws IOException {
+ // TODO: read config fields when added
+ }
+
+ @Override
+ public InstructionType type() {
+ return InstructionType.SETUP_PARTIAL_AGGREGATE;
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ // TODO: write config fields when added
+ }
+}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/SearchExecEngineProvider.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/SearchExecEngineProvider.java
index f16b8f36d9021..8edd8d0a71dc6 100644
--- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/SearchExecEngineProvider.java
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/SearchExecEngineProvider.java
@@ -9,8 +9,8 @@
package org.opensearch.analytics.spi;
import org.opensearch.analytics.backend.EngineResultStream;
-import org.opensearch.analytics.backend.ExecutionContext;
import org.opensearch.analytics.backend.SearchExecEngine;
+import org.opensearch.analytics.backend.ShardScanExecutionContext;
/**
* Execution engine factory for backend plugins.
@@ -23,6 +23,10 @@ public interface SearchExecEngineProvider {
/**
* Creates a search execution engine bound to the given execution context.
* The context carries the reader snapshot and task metadata.
+ * The backendContext carries backend-specific state produced by instruction handlers.
*/
- SearchExecEngine createSearchExecEngine(ExecutionContext ctx);
+ SearchExecEngine createSearchExecEngine(
+ ShardScanExecutionContext ctx,
+ BackendExecutionContext backendContext
+ );
}
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java
new file mode 100644
index 0000000000000..8000d34f68844
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java
@@ -0,0 +1,39 @@
+/*
+ * 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 base shard scan setup — reader acquisition, SessionContext creation,
+ * default table provider registration.
+ *
+ * @opensearch.internal
+ */
+public class ShardScanInstructionNode implements InstructionNode {
+
+ public ShardScanInstructionNode() {}
+
+ public ShardScanInstructionNode(StreamInput in) throws IOException {
+ // No fields to read
+ }
+
+ @Override
+ public InstructionType type() {
+ return InstructionType.SETUP_SHARD_SCAN;
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ // No fields to write
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
index 7c9a67827f326..ce180c060f563 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
@@ -412,6 +412,29 @@ pub unsafe extern "C" fn df_cache_manager_add_files(
Ok(0)
}
+// ---------------------------------------------------------------------------
+// SessionContext decomposition — instruction-based execution
+// ---------------------------------------------------------------------------
+
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_create_session_context(
+ shard_view_ptr: i64,
+ runtime_ptr: i64,
+ table_name_ptr: *const u8,
+ table_name_len: i64,
+ context_id: i64,
+) -> i64 {
+ let table_name = str_from_raw(table_name_ptr, table_name_len)
+ .map_err(|e| format!("df_create_session_context: {}", e))?;
+ let mgr = get_rt_manager()?;
+ mgr.io_runtime
+ .block_on(crate::session_context::create_session_context(
+ runtime_ptr, shard_view_ptr, table_name, context_id,
+ ))
+ .map_err(|e| e.to_string())
+}
+
#[ffm_safe]
#[no_mangle]
pub unsafe extern "C" fn df_cache_manager_remove_files(
@@ -526,3 +549,27 @@ pub unsafe extern "C" fn df_cache_manager_contains_by_type(
.ok_or_else(|| "df_cache_manager_contains_by_type: no cache manager configured".to_string())?;
Ok(if manager.contains_file_by_type(file_path, cache_type) { 1 } else { 0 })
}
+
+#[no_mangle]
+pub unsafe extern "C" fn df_close_session_context(ptr: i64) {
+ crate::session_context::close_session_context(ptr);
+}
+
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_execute_with_context(
+ session_ctx_ptr: i64,
+ plan_ptr: *const u8,
+ plan_len: i64,
+) -> i64 {
+ 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_ctx_ptr,
+ plan_bytes,
+ cpu_executor,
+ ))
+ .map_err(|e| e.to_string())
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
index d12fa24fa3acc..ee876888450c9 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
@@ -27,4 +27,5 @@ pub mod partition_stream;
pub mod query_executor;
pub mod query_memory_pool_tracker;
pub mod runtime_manager;
+pub mod session_context;
pub mod statistics_cache;
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs
index 766baffcc7afe..8ba9c93b3caea 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs
@@ -29,9 +29,15 @@ use substrait::proto::Plan;
use crate::cross_rt_stream::CrossRtStream;
use crate::executor::DedicatedExecutor;
use crate::api::DataFusionRuntime;
+use crate::session_context::SessionContextHandle;
/// Execute a vanilla parquet query: substrait plan → DataFusion → CrossRtStream.
/// File access goes through DataFusion's registered object store.
+///
+/// Deprecated: Production now uses the decomposed `create_session_context` +
+/// `execute_with_context` path (via `api::execute_query`).
+/// TODO: Remove this function and migrate benchmarks to the decomposed path.
+/// Retained only for benchmarks. TODO: migrate benchmarks and remove.
pub async fn execute_query(
table_path: ListingTableUrl,
object_metas: Arc>,
@@ -147,3 +153,35 @@ pub async fn execute_query(
Ok(Box::into_raw(Box::new(wrapped)) as i64)
}
+
+/// Executes a Substrait plan against a pre-configured SessionContext.
+/// Consumes the handle — SessionContext lifetime is tied to the returned stream.
+pub async unsafe fn execute_with_context(
+ session_ctx_ptr: i64,
+ plan_bytes: &[u8],
+ cpu_executor: DedicatedExecutor,
+) -> Result {
+ let handle = *Box::from_raw(session_ctx_ptr as *mut SessionContextHandle);
+
+ let substrait_plan = Plan::decode(plan_bytes).map_err(|e| {
+ DataFusionError::Execution(format!("Failed to decode Substrait: {}", e))
+ })?;
+
+ let logical_plan = from_substrait_plan(&handle.ctx.state(), &substrait_plan).await?;
+ let dataframe = handle.ctx.execute_logical_plan(logical_plan).await?;
+ let physical_plan = dataframe.create_physical_plan().await?;
+
+ let df_stream = execute_stream(physical_plan, handle.ctx.task_ctx()).map_err(|e| {
+ error!("execute_with_context: failed to create stream: {}", e);
+ e
+ })?;
+
+ let cross_rt_stream = CrossRtStream::new_with_df_error_stream(df_stream, cpu_executor);
+ let wrapped = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
+ cross_rt_stream.schema(),
+ cross_rt_stream,
+ );
+
+ let stream_handle = crate::api::QueryStreamHandle::new(wrapped, handle.query_context);
+ Ok(Box::into_raw(Box::new(stream_handle)) as i64)
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs
new file mode 100644
index 0000000000000..9de9caaa968a5
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs
@@ -0,0 +1,149 @@
+/*
+ * 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.
+ */
+
+//! SessionContext lifecycle for instruction-based execution.
+//!
+//! `create_session_context` creates a fully configured SessionContext with
+//! the default ListingTable registered. Called by ShardScanInstruction handler.
+
+use std::sync::Arc;
+
+use datafusion::{
+ common::DataFusionError,
+ datasource::file_format::parquet::ParquetFormat,
+ datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
+ execution::cache::cache_manager::CacheManagerConfig,
+ execution::cache::{CacheAccessor, DefaultListFilesCache},
+ execution::context::SessionContext,
+ execution::memory_pool::MemoryPool,
+ execution::runtime_env::RuntimeEnvBuilder,
+ execution::SessionStateBuilder,
+ prelude::*,
+};
+use log::error;
+use object_store::ObjectMeta;
+
+use crate::api::{DataFusionRuntime, ShardView};
+use crate::query_memory_pool_tracker::QueryTrackingContext;
+
+/// Opaque handle holding a configured SessionContext between FFM calls.
+pub struct SessionContextHandle {
+ pub ctx: SessionContext,
+ pub table_path: ListingTableUrl,
+ pub object_metas: Arc>,
+ pub query_context: QueryTrackingContext,
+}
+
+/// Creates a SessionContext with per-query RuntimeEnv and registers the default
+/// ListingTable provider for parquet scans.
+pub async unsafe fn create_session_context(
+ runtime_ptr: i64,
+ shard_view_ptr: i64,
+ table_name: &str,
+ context_id: i64,
+) -> Result {
+ let runtime = &*(runtime_ptr as *const DataFusionRuntime);
+ let shard_view = &*(shard_view_ptr as *const ShardView);
+
+ let global_pool = runtime.runtime_env.memory_pool.clone();
+ let query_context = QueryTrackingContext::new(context_id, global_pool);
+ let query_memory_pool = query_context
+ .memory_pool()
+ .map(|p| p as Arc);
+
+ let list_file_cache = Arc::new(DefaultListFilesCache::default());
+ list_file_cache.put(
+ &datafusion::execution::cache::TableScopedPath {
+ table: None,
+ path: shard_view.table_path.prefix().clone(),
+ },
+ shard_view.object_metas.clone(),
+ );
+
+ let mut runtime_env_builder = RuntimeEnvBuilder::from_runtime_env(&runtime.runtime_env)
+ .with_cache_manager(
+ CacheManagerConfig::default()
+ .with_list_files_cache(Some(list_file_cache))
+ .with_file_metadata_cache(Some(
+ runtime.runtime_env.cache_manager.get_file_metadata_cache(),
+ ))
+ .with_files_statistics_cache(
+ runtime.runtime_env.cache_manager.get_file_statistic_cache(),
+ ),
+ );
+
+ if let Some(pool) = query_memory_pool {
+ runtime_env_builder = runtime_env_builder.with_memory_pool(pool);
+ }
+
+ let runtime_env = runtime_env_builder.build().map_err(|e| {
+ error!("create_session_context: failed to build runtime env: {}", e);
+ e
+ })?;
+
+ let query_config = crate::datafusion_query_config::DatafusionQueryConfig::default();
+ let mut config = SessionConfig::new();
+ config.options_mut().execution.parquet.pushdown_filters = query_config.parquet_pushdown_filters;
+ config.options_mut().execution.target_partitions = query_config.target_partitions;
+ config.options_mut().execution.batch_size = query_config.batch_size;
+
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_runtime_env(Arc::from(runtime_env))
+ .with_default_features()
+ .build();
+
+ let ctx = SessionContext::new_with_state(state);
+
+ // Register default ListingTable for parquet scans
+ let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new()))
+ .with_file_extension(".parquet")
+ .with_collect_stat(true);
+
+ let resolved_schema = listing_options
+ .infer_schema(&ctx.state(), &shard_view.table_path)
+ .await
+ .map_err(|e| {
+ error!("create_session_context: failed to infer schema: {}", e);
+ e
+ })?;
+
+ let table_config = ListingTableConfig::new(shard_view.table_path.clone())
+ .with_listing_options(listing_options)
+ .with_schema(resolved_schema);
+
+ let provider = Arc::new(ListingTable::try_new(table_config).map_err(|e| {
+ error!("create_session_context: failed to create listing table: {}", e);
+ e
+ })?);
+
+ ctx.register_table(table_name, provider).map_err(|e| {
+ error!("create_session_context: failed to register table '{}': {}", table_name, e);
+ e
+ })?;
+
+ error!("create_session_context: successfully registered table '{}', table_name_len={}", table_name, table_name.len());
+
+ let handle = SessionContextHandle {
+ ctx,
+ table_path: shard_view.table_path.clone(),
+ object_metas: shard_view.object_metas.clone(),
+ query_context,
+ };
+ Ok(Box::into_raw(Box::new(handle)) as i64)
+}
+
+/// Closes a SessionContext handle without executing. Used for cleanup on failure.
+///
+/// # Safety
+/// `ptr` must be 0 or a valid pointer returned by `create_session_context`.
+pub unsafe fn close_session_context(ptr: i64) {
+ if ptr != 0 {
+ let _ = Box::from_raw(ptr as *mut SessionContextHandle);
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
index f5ed1de1033bd..b3730c709822a 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
@@ -17,6 +17,7 @@
import org.opensearch.analytics.spi.FieldType;
import org.opensearch.analytics.spi.FilterCapability;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
import org.opensearch.analytics.spi.ProjectCapability;
import org.opensearch.analytics.spi.ScalarFunction;
import org.opensearch.analytics.spi.ScalarFunctionAdapter;
@@ -153,7 +154,7 @@ public FragmentConvertor getFragmentConvertor() {
@Override
public SearchExecEngineProvider getSearchExecEngineProvider() {
- return ctx -> {
+ return (ctx, backendContext) -> {
DataFusionService dataFusionService = plugin.getDataFusionService();
if (dataFusionService == null) {
throw new IllegalStateException("DataFusionService not initialized — createComponents() may not have been called");
@@ -175,12 +176,21 @@ public SearchExecEngineProvider getSearchExecEngineProvider() {
throw new IllegalStateException("No DatafusionReader available in the acquired reader");
}
DatafusionContext context = new DatafusionContext(ctx.getTask(), dfReader, dataFusionService.getNativeRuntime());
+ if (backendContext != null) {
+ DataFusionSessionState sessionState = (DataFusionSessionState) backendContext;
+ context.setSessionContextHandle(sessionState.sessionContextHandle());
+ }
DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context);
engine.prepare(ctx);
return engine;
};
}
+ @Override
+ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
+ return new DataFusionInstructionHandlerFactory(plugin);
+ }
+
@Override
public ExchangeSinkProvider getExchangeSinkProvider() {
return ctx -> {
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java
new file mode 100644
index 0000000000000..88ec47e2da9f9
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java
@@ -0,0 +1,78 @@
+/*
+ * 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.be.datafusion;
+
+import org.opensearch.analytics.spi.DelegatedExpression;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
+import org.opensearch.analytics.spi.FinalAggregateInstructionNode;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
+
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * DataFusion backend's instruction handler factory.
+ *
+ * Coordinator side: creates typed instruction nodes for wire transport.
+ *
Data node side: creates handlers that call into Rust via FFM to configure the SessionContext.
+ */
+public class DataFusionInstructionHandlerFactory implements FragmentInstructionHandlerFactory {
+
+ private final DataFusionPlugin plugin;
+
+ public DataFusionInstructionHandlerFactory(DataFusionPlugin plugin) {
+ this.plugin = plugin;
+ }
+
+ // ── Coordinator: create instruction nodes ──
+
+ @Override
+ public Optional createShardScanNode() {
+ return Optional.of(new ShardScanInstructionNode());
+ }
+
+ @Override
+ public Optional createFilterDelegationNode(
+ FilterTreeShape treeShape,
+ int delegatedPredicateCount,
+ List delegatedExpressions
+ ) {
+ return Optional.of(new FilterDelegationInstructionNode(treeShape, delegatedPredicateCount, delegatedExpressions));
+ }
+
+ @Override
+ public Optional createPartialAggregateNode() {
+ // TODO: return Optional.of(...) once PartialAggregateInstructionHandler is implemented
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional createFinalAggregateNode() {
+ return Optional.of(new FinalAggregateInstructionNode());
+ }
+
+ // ── Data node: create handlers ──
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public FragmentInstructionHandler createHandler(InstructionNode node) {
+ if (node instanceof ShardScanInstructionNode) {
+ return new ShardScanInstructionHandler(plugin);
+ }
+ if (node instanceof FinalAggregateInstructionNode) {
+ return new FinalAggregateInstructionHandler();
+ }
+ // TODO: FilterDelegationInstructionHandler, PartialAggregateInstructionHandler
+ throw new UnsupportedOperationException("No handler for instruction type: " + node.type());
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java
new file mode 100644
index 0000000000000..edd48ddb11f22
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionSessionState.java
@@ -0,0 +1,19 @@
+/*
+ * 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.be.datafusion;
+
+import org.opensearch.analytics.spi.BackendExecutionContext;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
+
+/**
+ * Backend-specific execution context produced by ShardScanInstructionHandler,
+ * consumed by DatafusionSearcher at execute time.
+ */
+public record DataFusionSessionState(SessionContextHandle sessionContextHandle) implements BackendExecutionContext {
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java
index 527ead7df96a2..cd30c64e0758e 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionContext.java
@@ -8,6 +8,7 @@
package org.opensearch.be.datafusion;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
import org.opensearch.be.datafusion.nativelib.StreamHandle;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.search.SearchExecutionContext;
@@ -31,6 +32,7 @@ public class DatafusionContext implements SearchExecutionContext {
-
- private static final Logger logger = LogManager.getLogger(DatafusionSearchExecEngine.class);
+public class DatafusionSearchExecEngine implements SearchExecEngine {
private final DatafusionContext datafusionContext;
@@ -38,14 +34,14 @@ public DatafusionSearchExecEngine(DatafusionContext datafusionContext) {
}
@Override
- public void prepare(ExecutionContext requestContext) {
+ public void prepare(ShardScanExecutionContext requestContext) {
byte[] substraitBytes = requestContext.getFragmentBytes();
long contextId = datafusionContext.task() != null ? datafusionContext.task().getId() : 0L;
datafusionContext.setDatafusionQuery(new DatafusionQuery(requestContext.getTableName(), substraitBytes, contextId));
}
@Override
- public EngineResultStream execute(ExecutionContext requestContext) throws IOException {
+ public EngineResultStream execute(ShardScanExecutionContext requestContext) throws IOException {
BufferAllocator allocator = requestContext.getAllocator();
if (allocator == null) {
throw new IllegalStateException("ExecutionContext.allocator must be set by the caller before execute()");
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java
index 627ea17dd05e9..c8bb98991f10e 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSearcher.java
@@ -10,6 +10,7 @@
import org.opensearch.be.datafusion.nativelib.NativeBridge;
import org.opensearch.be.datafusion.nativelib.ReaderHandle;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
import org.opensearch.be.datafusion.nativelib.StreamHandle;
import org.opensearch.common.annotation.ExperimentalApi;
import org.opensearch.core.action.ActionListener;
@@ -45,6 +46,43 @@ public DatafusionSearcher(ReaderHandle readerHandle) {
@Override
public void search(DatafusionContext context) throws IOException {
+ SessionContextHandle sessionCtx = context.getSessionContextHandle();
+ if (sessionCtx != null) {
+ searchWithSessionContext(context, sessionCtx);
+ } else {
+ searchVanilla(context);
+ }
+ }
+
+ private void searchWithSessionContext(DatafusionContext context, SessionContextHandle sessionCtx) throws IOException {
+ DatafusionQuery query = context.getDatafusionQuery();
+ NativeRuntimeHandle runtimeHandle = context.getNativeRuntime();
+ CompletableFuture future = new CompletableFuture<>();
+ NativeBridge.executeWithContextAsync(sessionCtx.getPointer(), query.getSubstraitBytes(), new ActionListener<>() {
+ @Override
+ public void onResponse(Long streamPtr) {
+ future.complete(streamPtr);
+ }
+
+ @Override
+ public void onFailure(Exception exception) {
+ future.completeExceptionally(exception);
+ }
+ });
+ long streamPtr;
+ try {
+ streamPtr = future.join();
+ } catch (Exception exception) {
+ throw new IOException("Query execution with session context failed", exception);
+ }
+ // Rust consumed the session context — unregister from live handle set
+ sessionCtx.close();
+ context.setStreamHandle(new StreamHandle(streamPtr, runtimeHandle));
+ }
+
+ // TODO: Remove searchVanilla once all execution paths go through instruction handlers.
+ // Deprecated — retained only for tests that bypass AnalyticsSearchService.
+ private void searchVanilla(DatafusionContext context) throws IOException {
DatafusionQuery query = context.getDatafusionQuery();
if (query == null) {
throw new IllegalStateException("DatafusionQuery must be set before search");
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java
new file mode 100644
index 0000000000000..c3a9be1813c67
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java
@@ -0,0 +1,32 @@
+/*
+ * 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.be.datafusion;
+
+import org.opensearch.analytics.spi.BackendExecutionContext;
+import org.opensearch.analytics.spi.CommonExecutionContext;
+import org.opensearch.analytics.spi.FinalAggregateInstructionNode;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+
+/**
+ * Handles FinalAggregate instruction for coordinator-reduce stages.
+ * TODO: Configure SessionContext optimizer (disable CombinePartialFinalAggregate) for multi-shard aggregates.
+ */
+public class FinalAggregateInstructionHandler implements FragmentInstructionHandler {
+
+ @Override
+ public BackendExecutionContext apply(
+ FinalAggregateInstructionNode node,
+ CommonExecutionContext commonContext,
+ BackendExecutionContext backendContext
+ ) {
+ // TODO: Configure LocalSession optimizer settings for final aggregate execution.
+ // For now, the reduce path works without explicit configuration.
+ return backendContext;
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.java
new file mode 100644
index 0000000000000..2910499103bd5
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.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.be.datafusion;
+
+import org.opensearch.analytics.backend.ShardScanExecutionContext;
+import org.opensearch.analytics.spi.BackendExecutionContext;
+import org.opensearch.analytics.spi.CommonExecutionContext;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
+import org.opensearch.be.datafusion.nativelib.NativeBridge;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
+import org.opensearch.index.engine.dataformat.DataFormatRegistry;
+
+/**
+ * Handles ShardScan instruction: creates a SessionContext via FFM and registers
+ * the default ListingTable provider for parquet scans.
+ */
+public class ShardScanInstructionHandler implements FragmentInstructionHandler {
+
+ private final DataFusionPlugin plugin;
+
+ ShardScanInstructionHandler(DataFusionPlugin plugin) {
+ this.plugin = plugin;
+ }
+
+ @Override
+ public BackendExecutionContext apply(
+ ShardScanInstructionNode node,
+ CommonExecutionContext commonContext,
+ BackendExecutionContext backendContext
+ ) {
+ ShardScanExecutionContext context = (ShardScanExecutionContext) commonContext;
+ DataFusionService dataFusionService = plugin.getDataFusionService();
+ DataFormatRegistry registry = plugin.getDataFormatRegistry();
+
+ DatafusionReader dfReader = null;
+ for (String formatName : plugin.getSupportedFormats()) {
+ dfReader = context.getReader().getReader(registry.format(formatName), DatafusionReader.class);
+ if (dfReader != null) break;
+ }
+ if (dfReader == null) {
+ throw new IllegalStateException("No DatafusionReader available in the acquired reader");
+ }
+
+ long readerPtr = dfReader.getReaderHandle().getPointer();
+ long runtimePtr = dataFusionService.getNativeRuntime().get();
+ long contextId = context.getTask() != null ? context.getTask().getId() : 0L;
+
+ SessionContextHandle sessionCtxHandle = NativeBridge.createSessionContext(readerPtr, runtimePtr, context.getTableName(), contextId);
+
+ return new DataFusionSessionState(sessionCtxHandle);
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
index 387d18376cead..7afa5eb24feff 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
@@ -69,6 +69,9 @@ public final class NativeBridge {
private static final MethodHandle CACHE_MANAGER_GET_MEMORY_BY_TYPE;
private static final MethodHandle CACHE_MANAGER_GET_TOTAL_MEMORY;
private static final MethodHandle CACHE_MANAGER_CONTAINS_BY_TYPE;
+ private static final MethodHandle CREATE_SESSION_CONTEXT;
+ private static final MethodHandle CLOSE_SESSION_CONTEXT;
+ private static final MethodHandle EXECUTE_WITH_CONTEXT;
static {
SymbolLookup lib = NativeLibraryLoader.symbolLookup();
@@ -253,6 +256,19 @@ public final class NativeBridge {
)
);
+ // ── SessionContext decomposition bindings ──
+ CREATE_SESSION_CONTEXT = linker.downcallHandle(
+ lib.find("df_create_session_context").orElseThrow(),
+ FunctionDescriptor.of(
+ ValueLayout.JAVA_LONG,
+ ValueLayout.JAVA_LONG,
+ ValueLayout.JAVA_LONG,
+ ValueLayout.ADDRESS,
+ ValueLayout.JAVA_LONG,
+ ValueLayout.JAVA_LONG
+ )
+ );
+
// i64 df_cache_manager_add_files(runtime_ptr, files_ptr, files_len_ptr, files_count)
CACHE_MANAGER_ADD_FILES = linker.downcallHandle(
lib.find("df_cache_manager_add_files").orElseThrow(),
@@ -314,6 +330,16 @@ public final class NativeBridge {
// caller step required — as soon as this class is loaded, callbacks
// are installed and `df_execute_indexed_query` can dispatch into Java.
installFilterTreeCallbacks(linker);
+
+ CLOSE_SESSION_CONTEXT = linker.downcallHandle(
+ lib.find("df_close_session_context").orElseThrow(),
+ FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG)
+ );
+
+ EXECUTE_WITH_CONTEXT = linker.downcallHandle(
+ lib.find("df_execute_with_context").orElseThrow(),
+ FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG)
+ );
}
private NativeBridge() {}
@@ -650,6 +676,40 @@ public static long createCustomCacheManager() {
throw t instanceof RuntimeException ? (RuntimeException) t : new RuntimeException(t);
}
}
+ // ---- SessionContext decomposition ----
+
+ /**
+ * Creates a SessionContext with the default ListingTable registered.
+ * Returns a tracked handle consumed by {@link #executeWithContextAsync}.
+ */
+ public static SessionContextHandle createSessionContext(long readerPtr, long runtimePtr, String tableName, long contextId) {
+ NativeHandle.validatePointer(readerPtr, "reader");
+ NativeHandle.validatePointer(runtimePtr, "runtime");
+ try (var call = new NativeCall()) {
+ var table = call.str(tableName);
+ long ptr = call.invoke(CREATE_SESSION_CONTEXT, readerPtr, runtimePtr, table.segment(), table.len(), contextId);
+ return new SessionContextHandle(ptr);
+ }
+ }
+
+ /**
+ * Executes a Substrait plan against the configured SessionContext.
+ * Consumes the session context handle (freed internally when stream closes).
+ */
+ /** Frees a native SessionContext handle. Safe to call once. */
+ public static void closeSessionContext(long ptr) {
+ NativeCall.invokeVoid(CLOSE_SESSION_CONTEXT, ptr);
+ }
+
+ public static void executeWithContextAsync(long sessionCtxPtr, byte[] substraitPlan, ActionListener listener) {
+ NativeHandle.validatePointer(sessionCtxPtr, "sessionContext");
+ try (var call = new NativeCall()) {
+ long result = call.invoke(EXECUTE_WITH_CONTEXT, sessionCtxPtr, call.bytes(substraitPlan), (long) substraitPlan.length);
+ listener.onResponse(result);
+ } catch (Throwable throwable) {
+ listener.onFailure(throwable instanceof Exception ? (Exception) throwable : new RuntimeException(throwable));
+ }
+ }
public static void destroyCustomCacheManager(long ptr) {
NativeCall.invokeVoid(DESTROY_CUSTOM_CACHE_MANAGER, ptr);
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java
new file mode 100644
index 0000000000000..c26d1799611eb
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/SessionContextHandle.java
@@ -0,0 +1,38 @@
+/*
+ * 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.be.datafusion.nativelib;
+
+import org.opensearch.analytics.backend.jni.NativeHandle;
+
+/**
+ * Type-safe wrapper for a native SessionContext pointer returned by
+ * {@link NativeBridge#createSessionContext}. The Rust side consumes this
+ * handle when {@link NativeBridge#executeWithContextAsync} is called,
+ * so {@link #doClose()} is a no-op — the pointer is freed by Rust internally.
+ *
+ * This handle exists to participate in the {@link NativeHandle} live-pointer
+ * registry so that {@link NativeHandle#validatePointer} passes for FFM calls.
+ */
+public class SessionContextHandle extends NativeHandle {
+
+ public SessionContextHandle(long ptr) {
+ super(ptr);
+ }
+
+ @Override
+ protected void doClose() {
+ // TODO: Handle error-path cleanup. Currently Rust consumes the handle in
+ // execute_with_context (moves QueryTrackingContext into the stream). If execute
+ // fails or is never called, this handle leaks on the Rust side.
+ // Options: (a) AtomicBool 'consumed' flag on Rust handle — close_session_context
+ // checks flag before freeing, (b) don't consume in Rust and use no-op tracking
+ // on the stream, (c) markConsumed() on NativeHandle to skip doClose on happy path.
+ // See df_close_session_context FFM entry which exists but is not yet wired here.
+ }
+}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java
index 3ec318a4e2ae6..0e5c78087986f 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java
@@ -10,10 +10,13 @@
import org.opensearch.be.datafusion.nativelib.NativeBridge;
import org.opensearch.be.datafusion.nativelib.ReaderHandle;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
+import org.opensearch.core.action.ActionListener;
import org.opensearch.test.OpenSearchTestCase;
import java.nio.file.Files;
import java.nio.file.Path;
+import java.util.concurrent.CompletableFuture;
/**
* Smoke test for the DataFusion JNI bridge.
@@ -63,4 +66,55 @@ public void testReaderLifecycle() throws Exception {
NativeBridge.closeGlobalRuntime(runtimePtr);
}
+
+ public void testSessionContextCreationAndTableRegistration() throws Exception {
+ NativeBridge.initTokioRuntimeManager(2);
+ Path spillDir = createTempDir("datafusion-spill");
+ long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024);
+ NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr);
+
+ Path dataDir = createTempDir("datafusion-data");
+ Path testParquet = Path.of(getClass().getClassLoader().getResource("test.parquet").toURI());
+ Files.copy(testParquet, dataDir.resolve("test.parquet"));
+
+ ReaderHandle readerHandle = new ReaderHandle(dataDir.toString(), new String[] { "test.parquet" });
+
+ // Create session context with table registered
+ SessionContextHandle sessionCtx = NativeBridge.createSessionContext(
+ readerHandle.getPointer(),
+ runtimeHandle.get(),
+ "test_table",
+ 0L
+ );
+ assertTrue("SessionContext pointer should be non-zero", sessionCtx.getPointer() != 0);
+
+ // Execute a simple query to verify the session context is properly configured
+ byte[] substrait = NativeBridge.sqlToSubstrait(
+ readerHandle.getPointer(),
+ "test_table",
+ "SELECT message FROM test_table",
+ runtimeHandle.get()
+ );
+ CompletableFuture future = new CompletableFuture<>();
+ NativeBridge.executeWithContextAsync(sessionCtx.getPointer(), substrait, new ActionListener<>() {
+ @Override
+ public void onResponse(Long streamPtr) {
+ future.complete(streamPtr);
+ }
+
+ @Override
+ public void onFailure(Exception exception) {
+ future.completeExceptionally(exception);
+ }
+ });
+ long streamPtr = future.join();
+ assertTrue("Stream pointer should be non-zero", streamPtr != 0);
+
+ // Session context is consumed by execute — close the Java handle
+ sessionCtx.close();
+
+ NativeBridge.streamClose(streamPtr);
+ readerHandle.close();
+ runtimeHandle.close();
+ }
}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java
index 7135b988599e8..21fea0bc58494 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java
@@ -11,9 +11,10 @@
import org.apache.arrow.memory.RootAllocator;
import org.opensearch.analytics.backend.EngineResultBatch;
import org.opensearch.analytics.backend.EngineResultStream;
-import org.opensearch.analytics.backend.ExecutionContext;
+import org.opensearch.analytics.backend.ShardScanExecutionContext;
import org.opensearch.be.datafusion.nativelib.NativeBridge;
import org.opensearch.be.datafusion.nativelib.ReaderHandle;
+import org.opensearch.be.datafusion.nativelib.SessionContextHandle;
import org.opensearch.test.OpenSearchTestCase;
import java.nio.file.Files;
@@ -62,17 +63,16 @@ public void testEngineExecuteSelectAll() throws Exception {
runtimeHandle.get()
);
- // Build the plugin-level objects
DatafusionReader reader = createReader();
DatafusionContext context = new DatafusionContext(null, reader, runtimeHandle);
- context.setDatafusionQuery(new DatafusionQuery("test_table", substrait, 0L));
try (
RootAllocator alloc = new RootAllocator(Long.MAX_VALUE);
DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context)
) {
- ExecutionContext execCtx = new ExecutionContext("test_table", null, null);
+ ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context);
execCtx.setAllocator(alloc);
+ engine.prepare(execCtx);
try (EngineResultStream stream = engine.execute(execCtx)) {
List rows = collectRows(stream);
assertEquals(2, rows.size());
@@ -94,14 +94,14 @@ public void testEngineExecuteAggregation() throws Exception {
DatafusionReader reader = createReader();
DatafusionContext context = new DatafusionContext(null, reader, runtimeHandle);
- context.setDatafusionQuery(new DatafusionQuery("test_table", substrait, 0L));
try (
RootAllocator alloc = new RootAllocator(Long.MAX_VALUE);
DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context)
) {
- ExecutionContext execCtx = new ExecutionContext("test_table", null, null);
+ ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context);
execCtx.setAllocator(alloc);
+ engine.prepare(execCtx);
try (EngineResultStream stream = engine.execute(execCtx)) {
List rows = collectRows(stream);
assertEquals(1, rows.size());
@@ -120,14 +120,14 @@ public void testEngineExecuteFilter() throws Exception {
DatafusionReader reader = createReader();
DatafusionContext context = new DatafusionContext(null, reader, runtimeHandle);
- context.setDatafusionQuery(new DatafusionQuery("test_table", substrait, 0L));
try (
RootAllocator alloc = new RootAllocator(Long.MAX_VALUE);
DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context)
) {
- ExecutionContext execCtx = new ExecutionContext("test_table", null, null);
+ ShardScanExecutionContext execCtx = createExecutionContext("test_table", substrait, context);
execCtx.setAllocator(alloc);
+ engine.prepare(execCtx);
try (EngineResultStream stream = engine.execute(execCtx)) {
List rows = collectRows(stream);
assertEquals(1, rows.size());
@@ -137,10 +137,22 @@ public void testEngineExecuteFilter() throws Exception {
}
private DatafusionReader createReader() {
- // Wrap the raw pointer in a ReaderHandle via the existing native pointer
return new DatafusionReader(readerHandle.getPointer());
}
+ private ShardScanExecutionContext createExecutionContext(String tableName, byte[] substrait, DatafusionContext dfContext) {
+ ShardScanExecutionContext execCtx = new ShardScanExecutionContext(tableName, null, null);
+ execCtx.setFragmentBytes(substrait);
+ SessionContextHandle sessionCtxHandle = NativeBridge.createSessionContext(
+ readerHandle.getPointer(),
+ runtimeHandle.get(),
+ tableName,
+ 0L
+ );
+ dfContext.setSessionContextHandle(sessionCtxHandle);
+ return execCtx;
+ }
+
private List collectRows(EngineResultStream stream) {
List rows = new ArrayList<>();
Iterator it = stream.iterator();
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java
index 953f50afa1d09..bcbe2c3ef7c92 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java
@@ -38,15 +38,22 @@
import org.opensearch.analytics.planner.dag.StagePlan;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
import org.opensearch.analytics.spi.BackendCapabilityProvider;
+import org.opensearch.analytics.spi.DelegatedExpression;
import org.opensearch.analytics.spi.DelegatedPredicateFunction;
import org.opensearch.analytics.spi.DelegationType;
import org.opensearch.analytics.spi.EngineCapability;
import org.opensearch.analytics.spi.ExchangeSinkProvider;
import org.opensearch.analytics.spi.FieldType;
import org.opensearch.analytics.spi.FilterCapability;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
import org.opensearch.analytics.spi.ScalarFunction;
import org.opensearch.analytics.spi.ScanCapability;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
import org.opensearch.be.lucene.LuceneAnalyticsBackendPlugin;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.IndexMetadata;
@@ -69,6 +76,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
@@ -154,8 +162,8 @@ public void testMixedNativeAndDelegated() throws Exception {
);
StagePlan plan = runPipeline(condition);
- assertEquals("should have 1 delegated query", 1, plan.delegatedQueries().size());
- assertMatchQueryBuilder(plan.delegatedQueries(), "message", "hello world");
+ assertEquals("should have 1 delegated query", 1, plan.delegatedExpressions().size());
+ assertMatchQueryBuilder(plan.delegatedExpressions(), "message", "hello world");
SubstraitResult substrait = substraitResult(plan.convertedBytes());
logger.info("Substrait plan (mixed E2E):\n{}", substrait.plan());
@@ -165,7 +173,7 @@ public void testMixedNativeAndDelegated() throws Exception {
assertEquals("AND must have 2 arguments", 2, andFunc.getArgumentsCount());
// arg[1]: delegated_predicate(1) — annotation id=1 maps to MATCH 'hello world'
assertDelegatedPredicate(substrait.plan(), andFunc.getArguments(1).getValue(), 1);
- assertMatchQueryForAnnotation(plan.delegatedQueries(), 1, "message", "hello world");
+ assertMatchQueryForAnnotation(plan.delegatedExpressions(), 1, "message", "hello world");
}
/**
@@ -184,7 +192,7 @@ public void testComplexBooleanTree() throws Exception {
);
StagePlan plan = runPipeline(condition);
- assertEquals("should have 2 delegated queries", 2, plan.delegatedQueries().size());
+ assertEquals("should have 2 delegated queries", 2, plan.delegatedExpressions().size());
SubstraitResult substrait = substraitResult(plan.convertedBytes());
logger.info("Substrait plan (complex E2E):\n{}", substrait.plan());
@@ -203,14 +211,14 @@ public void testComplexBooleanTree() throws Exception {
// OR arg[0]: delegated_predicate(1) → MATCH 'hello'
assertDelegatedPredicate(substrait.plan(), orFunc.getArguments(0).getValue(), 1);
- assertMatchQueryForAnnotation(plan.delegatedQueries(), 1, "message", "hello");
+ assertMatchQueryForAnnotation(plan.delegatedExpressions(), 1, "message", "hello");
// OR arg[1]: NOT(delegated_predicate(2)) → MATCH 'goodbye'
Expression notExpr = orFunc.getArguments(1).getValue();
assertTrue("OR second arg must be scalar function", notExpr.hasScalarFunction());
assertEquals("not", resolveFunctionName(substrait.plan(), notExpr.getScalarFunction().getFunctionReference()));
assertDelegatedPredicate(substrait.plan(), notExpr.getScalarFunction().getArguments(0).getValue(), 2);
- assertMatchQueryForAnnotation(plan.delegatedQueries(), 2, "message", "goodbye");
+ assertMatchQueryForAnnotation(plan.delegatedExpressions(), 2, "message", "goodbye");
}
// ---- Pipeline ----
@@ -261,10 +269,10 @@ private RexNode makeMatch(int fieldIndex, String query) {
);
}
- private void assertMatchQueryBuilder(Map delegatedQueries, String expectedField, String expectedQuery)
+ private void assertMatchQueryBuilder(List delegatedExpressions, String expectedField, String expectedQuery)
throws IOException {
- for (byte[] queryBytes : delegatedQueries.values()) {
- try (StreamInput input = new NamedWriteableAwareStreamInput(StreamInput.wrap(queryBytes), WRITEABLE_REGISTRY)) {
+ for (DelegatedExpression expr : delegatedExpressions) {
+ try (StreamInput input = new NamedWriteableAwareStreamInput(StreamInput.wrap(expr.getExpressionBytes()), WRITEABLE_REGISTRY)) {
MatchQueryBuilder matchQuery = (MatchQueryBuilder) input.readNamedWriteable(QueryBuilder.class);
if (matchQuery.fieldName().equals(expectedField) && matchQuery.value().equals(expectedQuery)) {
return;
@@ -311,15 +319,20 @@ private void assertDelegatedPredicate(Plan plan, Expression expr, int expectedAn
/** Asserts the delegated query bytes for a specific annotation ID deserialize to the expected MatchQueryBuilder. */
private void assertMatchQueryForAnnotation(
- Map delegatedQueries,
+ List delegatedExpressions,
int annotationId,
String expectedField,
String expectedQuery
) throws IOException {
- assertTrue("annotation ID " + annotationId + " must be in delegatedQueries", delegatedQueries.containsKey(annotationId));
- try (
- StreamInput input = new NamedWriteableAwareStreamInput(StreamInput.wrap(delegatedQueries.get(annotationId)), WRITEABLE_REGISTRY)
- ) {
+ DelegatedExpression found = null;
+ for (DelegatedExpression expr : delegatedExpressions) {
+ if (expr.getAnnotationId() == annotationId) {
+ found = expr;
+ break;
+ }
+ }
+ assertNotNull("annotation ID " + annotationId + " must be in delegatedExpressions", found);
+ try (StreamInput input = new NamedWriteableAwareStreamInput(StreamInput.wrap(found.getExpressionBytes()), WRITEABLE_REGISTRY)) {
MatchQueryBuilder matchQuery = (MatchQueryBuilder) input.readNamedWriteable(QueryBuilder.class);
assertEquals("field name for annotation " + annotationId, expectedField, matchQuery.fieldName());
assertEquals("query text for annotation " + annotationId, expectedQuery, matchQuery.value());
@@ -437,5 +450,39 @@ public ExchangeSinkProvider getExchangeSinkProvider() {
public FragmentConvertor getFragmentConvertor() {
return convertor;
}
+
+ @Override
+ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
+ return new FragmentInstructionHandlerFactory() {
+ @Override
+ public Optional createShardScanNode() {
+ return Optional.of(new ShardScanInstructionNode());
+ }
+
+ @Override
+ public Optional createFilterDelegationNode(
+ FilterTreeShape treeShape,
+ int delegatedPredicateCount,
+ List delegatedExpressions
+ ) {
+ return Optional.of(new FilterDelegationInstructionNode(treeShape, delegatedPredicateCount, delegatedExpressions));
+ }
+
+ @Override
+ public Optional createPartialAggregateNode() {
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional createFinalAggregateNode() {
+ return Optional.empty();
+ }
+
+ @Override
+ public FragmentInstructionHandler createHandler(InstructionNode node) {
+ throw new UnsupportedOperationException("stub");
+ }
+ };
+ }
}
}
diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java
index 9f441489365a6..735ce496ba481 100644
--- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java
+++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java
@@ -37,14 +37,21 @@
import org.opensearch.analytics.planner.dag.StagePlan;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
import org.opensearch.analytics.spi.BackendCapabilityProvider;
+import org.opensearch.analytics.spi.DelegatedExpression;
import org.opensearch.analytics.spi.DelegationType;
import org.opensearch.analytics.spi.EngineCapability;
import org.opensearch.analytics.spi.ExchangeSinkProvider;
import org.opensearch.analytics.spi.FieldType;
import org.opensearch.analytics.spi.FilterCapability;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
import org.opensearch.analytics.spi.ScalarFunction;
import org.opensearch.analytics.spi.ScanCapability;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.MappingMetadata;
@@ -67,6 +74,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
@@ -140,11 +148,11 @@ public void testMatchPredicateDelegationEndToEnd() throws IOException {
StagePlan plan = leaf.getPlanAlternatives().getFirst();
// Verify delegation happened
- assertFalse("delegatedQueries should not be empty", plan.delegatedQueries().isEmpty());
- assertEquals("should have exactly one delegated query", 1, plan.delegatedQueries().size());
+ assertFalse("delegatedExpressions should not be empty", plan.delegatedExpressions().isEmpty());
+ assertEquals("should have exactly one delegated expression", 1, plan.delegatedExpressions().size());
// Deserialize and verify the MatchQueryBuilder
- byte[] queryBytes = plan.delegatedQueries().values().iterator().next();
+ byte[] queryBytes = plan.delegatedExpressions().getFirst().getExpressionBytes();
try (StreamInput input = new NamedWriteableAwareStreamInput(StreamInput.wrap(queryBytes), WRITEABLE_REGISTRY)) {
QueryBuilder deserialized = input.readNamedWriteable(QueryBuilder.class);
assertTrue("Should be MatchQueryBuilder", deserialized instanceof MatchQueryBuilder);
@@ -282,5 +290,39 @@ public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByte
}
};
}
+
+ @Override
+ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
+ return new FragmentInstructionHandlerFactory() {
+ @Override
+ public Optional createShardScanNode() {
+ return Optional.of(new ShardScanInstructionNode());
+ }
+
+ @Override
+ public Optional createFilterDelegationNode(
+ FilterTreeShape treeShape,
+ int delegatedPredicateCount,
+ List delegatedExpressions
+ ) {
+ return Optional.of(new FilterDelegationInstructionNode(treeShape, delegatedPredicateCount, delegatedExpressions));
+ }
+
+ @Override
+ public Optional createPartialAggregateNode() {
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional createFinalAggregateNode() {
+ return Optional.empty();
+ }
+
+ @Override
+ public FragmentInstructionHandler createHandler(InstructionNode node) {
+ throw new UnsupportedOperationException("mock");
+ }
+ };
+ }
}
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java
index 252347f2da2cf..804635187cf4d 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java
@@ -12,12 +12,16 @@
import org.opensearch.analytics.backend.AnalyticsOperationListener;
import org.opensearch.analytics.backend.EngineResultBatch;
import org.opensearch.analytics.backend.EngineResultStream;
-import org.opensearch.analytics.backend.ExecutionContext;
import org.opensearch.analytics.backend.SearchExecEngine;
+import org.opensearch.analytics.backend.ShardScanExecutionContext;
import org.opensearch.analytics.exec.action.FragmentExecutionRequest;
import org.opensearch.analytics.exec.action.FragmentExecutionResponse;
import org.opensearch.analytics.exec.task.AnalyticsShardTask;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
+import org.opensearch.analytics.spi.BackendExecutionContext;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
import org.opensearch.arrow.flight.transport.ArrowAllocatorProvider;
import org.opensearch.common.Nullable;
import org.opensearch.common.concurrent.GatedCloseable;
@@ -36,7 +40,7 @@
/**
* Data-node service that executes plan fragments against local shards.
* Acquires a reader from the shard's composite engine, builds an
- * {@link ExecutionContext}, and invokes the backend's {@link SearchExecEngine}
+ * {@link ShardScanExecutionContext}, and invokes the backend's {@link SearchExecEngine}
* to produce results.
*
* Does NOT hold {@code IndicesService} — receives an already-resolved
@@ -107,12 +111,24 @@ public FragmentResources executeFragmentStreaming(FragmentExecutionRequest reque
private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, Task task) throws IOException {
GatedCloseable gatedReader = resolved.readerProvider.acquireReader();
- SearchExecEngine engine = null;
+ SearchExecEngine engine = null;
EngineResultStream stream = null;
try {
- ExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, task);
+ ShardScanExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, task);
AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId());
- engine = backend.getSearchExecEngineProvider().createSearchExecEngine(ctx);
+
+ // Apply instruction handlers in order — each builds upon the previous handler's backend context
+ BackendExecutionContext backendContext = null;
+ List instructions = resolved.plan.getInstructions();
+ if (!instructions.isEmpty()) {
+ FragmentInstructionHandlerFactory factory = backend.getInstructionHandlerFactory();
+ for (InstructionNode node : instructions) {
+ FragmentInstructionHandler handler = factory.createHandler(node);
+ backendContext = handler.apply(node, ctx, backendContext);
+ }
+ }
+
+ engine = backend.getSearchExecEngineProvider().createSearchExecEngine(ctx, backendContext);
stream = engine.execute(ctx);
return new FragmentResources(gatedReader, engine, stream);
} catch (Exception e) {
@@ -158,13 +174,13 @@ private ResolvedFragment resolveFragment(FragmentExecutionRequest request, Index
return new ResolvedFragment(readerProvider, selectedPlan, request.getQueryId(), request.getStageId(), shardIdStr);
}
- private ExecutionContext buildContext(
+ private ShardScanExecutionContext buildContext(
FragmentExecutionRequest request,
Reader reader,
FragmentExecutionRequest.PlanAlternative plan,
Task task
) {
- ExecutionContext ctx = new ExecutionContext(request.getShardId().getIndexName(), task, reader);
+ ShardScanExecutionContext ctx = new ShardScanExecutionContext(request.getShardId().getIndexName(), task, reader);
ctx.setFragmentBytes(plan.getFragmentBytes());
ctx.setAllocator(allocator);
return ctx;
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java
index 9ad9170b4a812..0c61c3b85f79c 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java
@@ -9,8 +9,8 @@
package org.opensearch.analytics.exec;
import org.opensearch.analytics.backend.EngineResultStream;
-import org.opensearch.analytics.backend.ExecutionContext;
import org.opensearch.analytics.backend.SearchExecEngine;
+import org.opensearch.analytics.backend.ShardScanExecutionContext;
import org.opensearch.common.concurrent.GatedCloseable;
import org.opensearch.index.engine.exec.IndexReaderProvider.Reader;
@@ -23,12 +23,12 @@
public final class FragmentResources implements AutoCloseable {
private final GatedCloseable gatedReader;
- private final SearchExecEngine engine;
+ private final SearchExecEngine engine;
private final EngineResultStream stream;
public FragmentResources(
GatedCloseable gatedReader,
- SearchExecEngine engine,
+ SearchExecEngine engine,
EngineResultStream stream
) {
this.gatedReader = gatedReader;
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java
index ffc20d47fdb16..6bfb308081310 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java
@@ -11,6 +11,8 @@
import org.opensearch.action.ActionRequest;
import org.opensearch.action.ActionRequestValidationException;
import org.opensearch.analytics.exec.task.AnalyticsShardTask;
+import org.opensearch.analytics.spi.InstructionNode;
+import org.opensearch.analytics.spi.InstructionType;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.index.shard.ShardId;
@@ -97,28 +99,43 @@ public ActionRequestValidationException validate() {
}
/**
- * A single plan alternative: a backend ID paired with its serialized fragment bytes.
+ * A single plan alternative: a backend ID paired with its serialized fragment bytes
+ * and ordered instruction nodes for data-node execution.
* Produced by {@code FragmentConversionDriver.convertAll()} using the backend's
* {@code FragmentConvertor}.
*/
public static class PlanAlternative {
private final String backendId;
private final byte[] fragmentBytes;
+ private final List instructions;
- public PlanAlternative(String backendId, byte[] fragmentBytes) {
+ public PlanAlternative(String backendId, byte[] fragmentBytes, List instructions) {
this.backendId = backendId;
this.fragmentBytes = fragmentBytes;
+ this.instructions = instructions;
}
public PlanAlternative(StreamInput in) throws IOException {
this.backendId = in.readString();
byte[] bytes = in.readByteArray();
this.fragmentBytes = (bytes.length == 0) ? null : bytes;
+ int instructionCount = in.readVInt();
+ List nodes = new ArrayList<>(instructionCount);
+ for (int i = 0; i < instructionCount; i++) {
+ InstructionType type = in.readEnum(InstructionType.class);
+ nodes.add(type.readNode(in));
+ }
+ this.instructions = nodes;
}
public void writeTo(StreamOutput out) throws IOException {
out.writeString(backendId);
out.writeByteArray(fragmentBytes != null ? fragmentBytes : new byte[0]);
+ out.writeVInt(instructions.size());
+ for (InstructionNode node : instructions) {
+ out.writeEnum(node.type());
+ node.writeTo(out);
+ }
}
public String getBackendId() {
@@ -128,5 +145,9 @@ public String getBackendId() {
public byte[] getFragmentBytes() {
return fragmentBytes;
}
+
+ public List getInstructions() {
+ return instructions;
+ }
}
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java
index 13aa77c61fb62..d9205260000ef 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java
@@ -11,9 +11,13 @@
import org.opensearch.analytics.exec.QueryContext;
import org.opensearch.analytics.planner.dag.Stage;
import org.opensearch.analytics.planner.dag.StageExecutionType;
+import org.opensearch.analytics.spi.BackendExecutionContext;
import org.opensearch.analytics.spi.ExchangeSink;
import org.opensearch.analytics.spi.ExchangeSinkContext;
import org.opensearch.analytics.spi.ExchangeSinkProvider;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
import java.util.ArrayList;
import java.util.List;
@@ -47,6 +51,22 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte
buildChildInputs(stage),
sink
);
+
+ // Apply instruction handlers for the reduce stage.
+ // Unlike AnalyticsSearchService (shard path) which resolves the factory from its
+ // local backends map, the coordinator-reduce path has no backends map — the factory
+ // is stored on the Stage during FragmentConversionDriver.convertAll (root stage only,
+ // no serialization needed since reduce executes locally at the coordinator).
+ // TODO: find a cleaner way to provide the factory without storing it on Stage.
+ FragmentInstructionHandlerFactory factory = stage.getInstructionHandlerFactory();
+ if (factory != null) {
+ BackendExecutionContext backendContext = null;
+ for (InstructionNode node : stage.getPlanAlternatives().getFirst().instructions()) {
+ FragmentInstructionHandler handler = factory.createHandler(node);
+ backendContext = handler.apply(node, context, backendContext);
+ }
+ }
+
ExchangeSink backendSink;
try {
backendSink = provider.createSink(context);
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java
index 701f0d2871e54..fdabe2581f423 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java
@@ -76,7 +76,7 @@ public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryConte
private static List buildPlanAlternatives(Stage stage) {
List alternatives = new ArrayList<>();
for (StagePlan plan : stage.getPlanAlternatives()) {
- alternatives.add(new FragmentExecutionRequest.PlanAlternative(plan.backendId(), plan.convertedBytes()));
+ alternatives.add(new FragmentExecutionRequest.PlanAlternative(plan.backendId(), plan.convertedBytes(), plan.instructions()));
}
return alternatives;
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java
index 5482d7dfe6b80..01474ba800efe 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java
@@ -98,6 +98,21 @@ public CapabilityRegistry(
for (DelegationType type : caps.supportedDelegations()) {
delegationSupporters.computeIfAbsent(type, k -> new ArrayList<>()).add(name);
}
+ // Validate: if a backend supports FILTER delegation (i.e., it drives the tree walk),
+ // it must provide a FragmentInstructionHandlerFactory for instruction-based execution.
+ if (caps.supportedDelegations().contains(DelegationType.FILTER)) {
+ try {
+ backend.getInstructionHandlerFactory();
+ } catch (UnsupportedOperationException exception) {
+ throw new IllegalStateException(
+ "Backend ["
+ + name
+ + "] declares supportedDelegations(FILTER) but does not implement"
+ + " getInstructionHandlerFactory(). A driving backend must provide an instruction"
+ + " handler factory to configure delegation at the data node."
+ );
+ }
+ }
for (DelegationType type : caps.acceptedDelegations()) {
delegationAcceptors.computeIfAbsent(type, k -> new ArrayList<>()).add(name);
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java
index 3e93f4eca9ee0..06cb3e725caa8 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java
@@ -19,7 +19,6 @@
import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer;
import org.opensearch.analytics.planner.rel.OpenSearchFilter;
import org.opensearch.analytics.planner.rel.OpenSearchProject;
-import org.opensearch.analytics.planner.rel.OpenSearchRelNode;
import org.opensearch.analytics.planner.rel.OpenSearchSort;
import org.opensearch.analytics.planner.rel.OpenSearchTableScan;
import org.opensearch.analytics.planner.rel.OpenSearchUnion;
@@ -112,29 +111,20 @@ private static RelTraitSet rebuildTraits(RelNode node, RelOptCluster newCluster,
}
/**
- * Extracts the single backend from the leaf operator in a resolved fragment.
- * After resolution, every operator has exactly one viable backend. Throws if
- * the leaf has more than one (indicates resolution didn't complete).
+ * Finds the first node of the given type in the fragment's single-input chain.
+ * Returns {@code null} if not found.
+ *
+ * TODO: migrate existing findLeaf/findFilter usages in FragmentConversionDriver to use this.
*/
- public static String extractLeafBackendFromResolvedFragment(RelNode node) {
- if (node.getInputs().isEmpty()) {
- if (node instanceof OpenSearchRelNode leafNode) {
- List backends = leafNode.getViableBackends();
- if (backends.size() != 1) {
- throw new IllegalStateException(
- "Expected exactly 1 viable backend on resolved leaf [" + node.getClass().getSimpleName() + "], got " + backends
- );
- }
- return backends.getFirst();
- }
- throw new IllegalStateException("Leaf node [" + node.getClass().getSimpleName() + "] is not an OpenSearchRelNode");
+ @SuppressWarnings("unchecked")
+ public static T findNode(RelNode node, Class type) {
+ if (type.isInstance(node)) {
+ return (T) node;
}
- for (RelNode input : node.getInputs()) {
- String backend = extractLeafBackendFromResolvedFragment(input);
- if (backend != null) {
- return backend;
- }
+ if (!node.getInputs().isEmpty()) {
+ return findNode(node.getInputs().getFirst(), type);
}
return null;
}
+
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriver.java
new file mode 100644
index 0000000000000..55123d261f56c
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriver.java
@@ -0,0 +1,76 @@
+/*
+ * 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.planner.dag;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlKind;
+import org.opensearch.analytics.planner.rel.AnnotatedPredicate;
+import org.opensearch.analytics.planner.rel.OpenSearchFilter;
+import org.opensearch.analytics.spi.FilterTreeShape;
+
+/**
+ * Derives {@link FilterTreeShape} from a filter condition while annotations are intact.
+ * Must be called before stripping removes the annotations.
+ *
+ * Single-pass walk: determines both whether delegation exists and whether the tree
+ * is mixed (delegated + driving-backend predicates interleaved under OR/NOT).
+ *
+ * @opensearch.internal
+ */
+final class FilterTreeShapeDeriver {
+
+ private FilterTreeShapeDeriver() {}
+
+ /**
+ * Derives the filter tree shape from the filter's condition.
+ *
+ * @param filter the OpenSearchFilter with annotations intact
+ * @param drivingBackendId the filter operator's resolved backend
+ * @return the tree shape, or {@code null} if no delegated annotations exist
+ */
+ static FilterTreeShape derive(OpenSearchFilter filter, String drivingBackendId) {
+ Result result = walk(filter.getCondition(), drivingBackendId);
+ if (!result.hasDelegated) {
+ return FilterTreeShape.NO_DELEGATION;
+ }
+ return result.hasMixed ? FilterTreeShape.INTERLEAVED_BOOLEAN_EXPRESSION : FilterTreeShape.CONJUNCTIVE;
+ }
+
+ private static Result walk(RexNode node, String drivingBackendId) {
+ if (node instanceof AnnotatedPredicate predicate) {
+ boolean isDelegated = !predicate.getViableBackends().getFirst().equals(drivingBackendId);
+ return new Result(isDelegated, false, !isDelegated);
+ }
+ if (node instanceof RexCall call) {
+ boolean isOrNot = call.getKind() == SqlKind.OR || call.getKind() == SqlKind.NOT;
+
+ boolean hasDelegated = false;
+ boolean hasDrivingBackend = false;
+ boolean hasMixed = false;
+
+ for (RexNode operand : call.getOperands()) {
+ Result childResult = walk(operand, drivingBackendId);
+ hasDelegated |= childResult.hasDelegated;
+ hasDrivingBackend |= childResult.hasDrivingBackend;
+ hasMixed |= childResult.hasMixed;
+ }
+
+ if (isOrNot && hasDelegated && hasDrivingBackend) {
+ hasMixed = true;
+ }
+
+ return new Result(hasDelegated, hasMixed, hasDrivingBackend);
+ }
+ return new Result(false, false, false);
+ }
+
+ private record Result(boolean hasDelegated, boolean hasMixed, boolean hasDrivingBackend) {
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java
index 3a346a8ff6efa..c97489e9d239b 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java
@@ -16,23 +16,28 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.analytics.planner.CapabilityRegistry;
+import org.opensearch.analytics.planner.RelNodeUtils;
import org.opensearch.analytics.planner.rel.AggregateMode;
import org.opensearch.analytics.planner.rel.OpenSearchAggregate;
import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer;
+import org.opensearch.analytics.planner.rel.OpenSearchFilter;
import org.opensearch.analytics.planner.rel.OpenSearchRelNode;
import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan;
import org.opensearch.analytics.planner.rel.OpenSearchTableScan;
import org.opensearch.analytics.planner.rel.OperatorAnnotation;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
+import org.opensearch.analytics.spi.DelegatedExpression;
import org.opensearch.analytics.spi.DelegatedPredicateSerializer;
import org.opensearch.analytics.spi.FieldStorageInfo;
+import org.opensearch.analytics.spi.FilterTreeShape;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
import org.opensearch.analytics.spi.ScalarFunction;
import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.LinkedList;
import java.util.List;
-import java.util.Map;
import java.util.function.Function;
/**
@@ -70,6 +75,12 @@ private FragmentConversionDriver() {}
*/
public static void convertAll(QueryDAG dag, CapabilityRegistry registry) {
convertStage(dag.rootStage(), registry);
+ // Root stage executes locally at coordinator — store factory for instruction dispatch.
+ Stage root = dag.rootStage();
+ if (root.getExchangeSinkProvider() != null && !root.getPlanAlternatives().isEmpty()) {
+ AnalyticsSearchBackendPlugin backend = registry.getBackend(root.getPlanAlternatives().getFirst().backendId());
+ root.setInstructionHandlerFactory(backend.getInstructionHandlerFactory());
+ }
}
private static void convertStage(Stage stage, CapabilityRegistry registry) {
@@ -80,11 +91,54 @@ private static void convertStage(Stage stage, CapabilityRegistry registry) {
for (StagePlan plan : stage.getPlanAlternatives()) {
AnalyticsSearchBackendPlugin backend = registry.getBackend(plan.backendId());
FragmentConvertor convertor = backend.getFragmentConvertor();
+
+ // Derive filter tree shape BEFORE stripping (annotations must be intact)
+ OpenSearchFilter filter = RelNodeUtils.findNode(plan.resolvedFragment(), OpenSearchFilter.class);
+ FilterTreeShape treeShape = filter != null
+ ? FilterTreeShapeDeriver.derive(filter, plan.backendId())
+ : FilterTreeShape.NO_DELEGATION;
+
IntraOperatorDelegationBytes delegationBytes = new IntraOperatorDelegationBytes(registry);
byte[] bytes = convert(plan.resolvedFragment(), convertor, delegationBytes);
- converted.add(plan.withConvertedBytes(bytes, delegationBytes.getResult()));
+
+ // Assemble instruction list
+ List instructions = assembleInstructions(backend, plan, treeShape, delegationBytes);
+
+ converted.add(plan.withConvertedBytes(bytes, delegationBytes.getResult()).withInstructions(instructions));
}
stage.setPlanAlternatives(converted);
+ // Store factory on coordinator-reduce stages (local execution, no serialization needed).
+ // Shard stages get the factory from the local backend plugin at the data node.
+ if (stage.getExchangeSinkProvider() != null && !converted.isEmpty()) {
+ AnalyticsSearchBackendPlugin backend = registry.getBackend(converted.getFirst().backendId());
+ stage.setInstructionHandlerFactory(backend.getInstructionHandlerFactory());
+ }
+ }
+
+ private static List assembleInstructions(
+ AnalyticsSearchBackendPlugin backend,
+ StagePlan plan,
+ FilterTreeShape treeShape,
+ IntraOperatorDelegationBytes delegationBytes
+ ) {
+ FragmentInstructionHandlerFactory factory = backend.getInstructionHandlerFactory();
+ LinkedList instructions = new LinkedList<>();
+ RelNode leaf = findLeaf(plan.resolvedFragment());
+
+ if (leaf instanceof OpenSearchTableScan) {
+ factory.createShardScanNode().ifPresent(instructions::add);
+ List delegated = delegationBytes.getResult();
+ if (!delegated.isEmpty()) {
+ factory.createFilterDelegationNode(treeShape, delegated.size(), delegated).ifPresent(instructions::add);
+ }
+ if (plan.resolvedFragment() instanceof OpenSearchAggregate agg && agg.getMode() == AggregateMode.PARTIAL) {
+ factory.createPartialAggregateNode().ifPresent(instructions::add);
+ }
+ } else if (leaf instanceof OpenSearchStageInputScan) {
+ factory.createFinalAggregateNode().ifPresent(instructions::add);
+ }
+
+ return instructions;
}
/**
@@ -93,7 +147,7 @@ private static void convertStage(Stage stage, CapabilityRegistry registry) {
*/
static final class IntraOperatorDelegationBytes {
private final CapabilityRegistry registry;
- private Map result;
+ private List delegatedExpressions;
IntraOperatorDelegationBytes(CapabilityRegistry registry) {
this.registry = registry;
@@ -140,16 +194,16 @@ Function resolverFor(OpenSearchRelNode operator, Re
annotationBackend,
serialized.length
);
- if (result == null) {
- result = new HashMap<>();
+ if (delegatedExpressions == null) {
+ delegatedExpressions = new ArrayList<>();
}
- result.put(annotation.getAnnotationId(), serialized);
+ delegatedExpressions.add(new DelegatedExpression(annotation.getAnnotationId(), annotationBackend, serialized));
return annotation.makePlaceholder(rexBuilder);
};
}
- Map getResult() {
- return result != null ? result : Map.of();
+ List getResult() {
+ return delegatedExpressions != null ? delegatedExpressions : List.of();
}
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java
index 410d657a691af..61e5668b5dda9 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java
@@ -10,6 +10,7 @@
import org.apache.calcite.rel.RelNode;
import org.opensearch.analytics.spi.ExchangeSinkProvider;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
import org.opensearch.common.Nullable;
import java.util.List;
@@ -45,6 +46,7 @@ public class Stage {
private final TargetResolver targetResolver;
private final StageExecutionType executionType;
private List planAlternatives;
+ private FragmentInstructionHandlerFactory instructionHandlerFactory;
public Stage(
int stageId,
@@ -118,6 +120,14 @@ public void setPlanAlternatives(List planAlternatives) {
this.planAlternatives = planAlternatives;
}
+ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
+ return instructionHandlerFactory;
+ }
+
+ public void setInstructionHandlerFactory(FragmentInstructionHandlerFactory instructionHandlerFactory) {
+ this.instructionHandlerFactory = instructionHandlerFactory;
+ }
+
private StageExecutionType setStageExecutionType(ExchangeSinkProvider exchangeSinkProvider, TargetResolver targetResolver) {
if (targetResolver != null) {
return StageExecutionType.SHARD_FRAGMENT;
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java
index 515a1d0537231..afa941ccaa5c3 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java
@@ -9,9 +9,11 @@
package org.opensearch.analytics.planner.dag;
import org.apache.calcite.rel.RelNode;
+import org.opensearch.analytics.spi.DelegatedExpression;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.InstructionNode;
-import java.util.Map;
+import java.util.List;
/**
* A single plan alternative for a {@link Stage}. Contains a resolved fragment
@@ -19,21 +21,28 @@
* are narrowed to exactly one backend, plus the converted bytes produced by
* the backend's {@link FragmentConvertor}.
*
- * @param resolvedFragment fragment with all viableBackends narrowed to single choices
- * @param backendId the primary backend for this plan
- * @param convertedBytes backend-specific serialized plan bytes (null before conversion)
- * @param delegatedQueries serialized delegated queries keyed by annotationId (empty if no delegation)
+ * @param resolvedFragment fragment with all viableBackends narrowed to single choices
+ * @param backendId the primary backend for this plan
+ * @param convertedBytes backend-specific serialized plan bytes (null before conversion)
+ * @param delegatedExpressions serialized delegated expressions (empty if no delegation)
+ * @param instructions ordered instruction nodes for data-node execution (empty before resolution)
* @opensearch.internal
*/
-public record StagePlan(RelNode resolvedFragment, String backendId, byte[] convertedBytes, Map delegatedQueries) {
+public record StagePlan(RelNode resolvedFragment, String backendId, byte[] convertedBytes, List delegatedExpressions,
+ List instructions) {
/** Creates a StagePlan before conversion (bytes not yet available). */
public StagePlan(RelNode resolvedFragment, String backendId) {
- this(resolvedFragment, backendId, null, Map.of());
+ this(resolvedFragment, backendId, null, List.of(), List.of());
}
- /** Returns a copy with converted bytes and delegated queries populated. */
- public StagePlan withConvertedBytes(byte[] bytes, Map delegatedQueries) {
- return new StagePlan(resolvedFragment, backendId, bytes, delegatedQueries);
+ /** Returns a copy with converted bytes and delegated expressions populated. */
+ public StagePlan withConvertedBytes(byte[] bytes, List delegatedExpressions) {
+ return new StagePlan(resolvedFragment, backendId, bytes, delegatedExpressions, List.of());
+ }
+
+ /** Returns a copy with instructions populated. */
+ public StagePlan withInstructions(List instructions) {
+ return new StagePlan(resolvedFragment, backendId, convertedBytes, delegatedExpressions, instructions);
}
}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/action/PlanAlternativeSerializationTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/action/PlanAlternativeSerializationTests.java
new file mode 100644
index 0000000000000..bc6d36bd8a62d
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/action/PlanAlternativeSerializationTests.java
@@ -0,0 +1,107 @@
+/*
+ * 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.exec.action;
+
+import org.opensearch.analytics.spi.DelegatedExpression;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
+import org.opensearch.analytics.spi.FinalAggregateInstructionNode;
+import org.opensearch.analytics.spi.InstructionNode;
+import org.opensearch.analytics.spi.InstructionType;
+import org.opensearch.analytics.spi.PartialAggregateInstructionNode;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Wire serialization round-trip tests for {@link FragmentExecutionRequest.PlanAlternative}.
+ */
+public class PlanAlternativeSerializationTests extends OpenSearchTestCase {
+
+ public void testRoundTripWithShardScanOnly() throws IOException {
+ List instructions = List.of(new ShardScanInstructionNode());
+ FragmentExecutionRequest.PlanAlternative original = new FragmentExecutionRequest.PlanAlternative(
+ "datafusion",
+ new byte[] { 1, 2, 3 },
+ instructions
+ );
+
+ FragmentExecutionRequest.PlanAlternative deserialized = roundTrip(original);
+
+ assertEquals("datafusion", deserialized.getBackendId());
+ assertArrayEquals(new byte[] { 1, 2, 3 }, deserialized.getFragmentBytes());
+ assertEquals(1, deserialized.getInstructions().size());
+ assertEquals(InstructionType.SETUP_SHARD_SCAN, deserialized.getInstructions().get(0).type());
+ }
+
+ public void testRoundTripWithFilterDelegation() throws IOException {
+ List expressions = List.of(
+ new DelegatedExpression(1, "lucene", new byte[] { 10, 20 }),
+ new DelegatedExpression(2, "lucene", new byte[] { 30, 40 })
+ );
+ FilterDelegationInstructionNode filterNode = new FilterDelegationInstructionNode(FilterTreeShape.CONJUNCTIVE, 2, expressions);
+ List instructions = List.of(new ShardScanInstructionNode(), filterNode);
+ FragmentExecutionRequest.PlanAlternative original = new FragmentExecutionRequest.PlanAlternative(
+ "datafusion",
+ new byte[] { 5, 6 },
+ instructions
+ );
+
+ FragmentExecutionRequest.PlanAlternative deserialized = roundTrip(original);
+
+ assertEquals(2, deserialized.getInstructions().size());
+ assertEquals(InstructionType.SETUP_SHARD_SCAN, deserialized.getInstructions().get(0).type());
+ assertEquals(InstructionType.SETUP_FILTER_DELEGATION_FOR_INDEX, deserialized.getInstructions().get(1).type());
+
+ FilterDelegationInstructionNode deserializedFilter = (FilterDelegationInstructionNode) deserialized.getInstructions().get(1);
+ assertEquals(FilterTreeShape.CONJUNCTIVE, deserializedFilter.getTreeShape());
+ assertEquals(2, deserializedFilter.getDelegatedPredicateCount());
+ assertEquals(2, deserializedFilter.getDelegatedQueries().size());
+ assertEquals(1, deserializedFilter.getDelegatedQueries().get(0).getAnnotationId());
+ assertEquals("lucene", deserializedFilter.getDelegatedQueries().get(0).getAcceptingBackendId());
+ assertArrayEquals(new byte[] { 10, 20 }, deserializedFilter.getDelegatedQueries().get(0).getExpressionBytes());
+ }
+
+ public void testRoundTripWithAllTypes() throws IOException {
+ List instructions = List.of(
+ new ShardScanInstructionNode(),
+ new FilterDelegationInstructionNode(
+ FilterTreeShape.INTERLEAVED_BOOLEAN_EXPRESSION,
+ 1,
+ List.of(new DelegatedExpression(3, "lucene", new byte[] { 99 }))
+ ),
+ new PartialAggregateInstructionNode(),
+ new FinalAggregateInstructionNode()
+ );
+ FragmentExecutionRequest.PlanAlternative original = new FragmentExecutionRequest.PlanAlternative(
+ "datafusion",
+ new byte[] { 7 },
+ instructions
+ );
+
+ FragmentExecutionRequest.PlanAlternative deserialized = roundTrip(original);
+
+ assertEquals(4, deserialized.getInstructions().size());
+ assertEquals(InstructionType.SETUP_SHARD_SCAN, deserialized.getInstructions().get(0).type());
+ assertEquals(InstructionType.SETUP_FILTER_DELEGATION_FOR_INDEX, deserialized.getInstructions().get(1).type());
+ assertEquals(InstructionType.SETUP_PARTIAL_AGGREGATE, deserialized.getInstructions().get(2).type());
+ assertEquals(InstructionType.SETUP_FINAL_AGGREGATE, deserialized.getInstructions().get(3).type());
+ }
+
+ private FragmentExecutionRequest.PlanAlternative roundTrip(FragmentExecutionRequest.PlanAlternative original) throws IOException {
+ BytesStreamOutput out = new BytesStreamOutput();
+ original.writeTo(out);
+ StreamInput in = out.bytes().streamInput();
+ return new FragmentExecutionRequest.PlanAlternative(in);
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java
index d1715c2ae3de4..aa8b9a054faf2 100644
--- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java
@@ -177,7 +177,7 @@ private ShardFragmentStage
"test-query",
0,
target.shardId(),
- List.of(new FragmentExecutionRequest.PlanAlternative("test-backend", new byte[0]))
+ List.of(new FragmentExecutionRequest.PlanAlternative("test-backend", new byte[0], List.of()))
);
return new ShardFragmentStageExecution(stage, config, sink, clusterService, requestBuilder, dispatcher, codec);
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java
index fa41b4fc9d5a7..a26f054ff34d0 100644
--- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java
@@ -28,7 +28,9 @@
import org.opensearch.analytics.planner.rel.OpenSearchFilter;
import org.opensearch.analytics.planner.rel.OpenSearchTableScan;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
+import org.opensearch.analytics.spi.BackendCapabilityProvider;
import org.opensearch.analytics.spi.DelegationType;
+import org.opensearch.analytics.spi.EngineCapability;
import java.util.List;
import java.util.Map;
@@ -341,4 +343,35 @@ protected Set acceptedDelegations() {
};
return List.of(df, lucene);
}
+
+ public void testBackendWithFilterDelegationButNoFactory_throws() {
+ AnalyticsSearchBackendPlugin badBackend = new AnalyticsSearchBackendPlugin() {
+ @Override
+ public String name() {
+ return "bad-backend";
+ }
+
+ @Override
+ public BackendCapabilityProvider getCapabilityProvider() {
+ return new BackendCapabilityProvider() {
+ @Override
+ public Set supportedEngineCapabilities() {
+ return Set.of();
+ }
+
+ @Override
+ public Set supportedDelegations() {
+ return Set.of(DelegationType.FILTER);
+ }
+ };
+ }
+ };
+
+ IllegalStateException exception = expectThrows(
+ IllegalStateException.class,
+ () -> new CapabilityRegistry(List.of(badBackend), idx -> null)
+ );
+ assertTrue(exception.getMessage().contains("bad-backend"));
+ assertTrue(exception.getMessage().contains("getInstructionHandlerFactory"));
+ }
}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java
index 33b70bd35ca5f..c378b00c6e254 100644
--- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java
@@ -11,16 +11,27 @@
import org.opensearch.analytics.spi.AggregateCapability;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
import org.opensearch.analytics.spi.BackendCapabilityProvider;
+import org.opensearch.analytics.spi.DelegatedExpression;
import org.opensearch.analytics.spi.DelegatedPredicateSerializer;
import org.opensearch.analytics.spi.DelegationType;
import org.opensearch.analytics.spi.EngineCapability;
import org.opensearch.analytics.spi.FilterCapability;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
+import org.opensearch.analytics.spi.FinalAggregateInstructionNode;
+import org.opensearch.analytics.spi.FragmentInstructionHandler;
+import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory;
+import org.opensearch.analytics.spi.InstructionNode;
+import org.opensearch.analytics.spi.PartialAggregateInstructionNode;
import org.opensearch.analytics.spi.ProjectCapability;
import org.opensearch.analytics.spi.ScalarFunction;
import org.opensearch.analytics.spi.ScalarFunctionAdapter;
import org.opensearch.analytics.spi.ScanCapability;
+import org.opensearch.analytics.spi.ShardScanInstructionNode;
+import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
/**
@@ -119,4 +130,38 @@ protected Map scalarFunctionAdapters() {
protected Map delegatedPredicateSerializers() {
return Map.of();
}
+
+ @Override
+ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() {
+ return new FragmentInstructionHandlerFactory() {
+ @Override
+ public Optional createShardScanNode() {
+ return Optional.of(new ShardScanInstructionNode());
+ }
+
+ @Override
+ public Optional createFilterDelegationNode(
+ FilterTreeShape treeShape,
+ int delegatedPredicateCount,
+ List delegatedExpressions
+ ) {
+ return Optional.of(new FilterDelegationInstructionNode(treeShape, delegatedPredicateCount, delegatedExpressions));
+ }
+
+ @Override
+ public Optional createPartialAggregateNode() {
+ return Optional.of(new PartialAggregateInstructionNode());
+ }
+
+ @Override
+ public Optional createFinalAggregateNode() {
+ return Optional.of(new FinalAggregateInstructionNode());
+ }
+
+ @Override
+ public FragmentInstructionHandler createHandler(InstructionNode node) {
+ throw new UnsupportedOperationException("Mock backend does not execute instructions");
+ }
+ };
+ }
}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriverTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriverTests.java
new file mode 100644
index 0000000000000..8930b1043b9a7
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FilterTreeShapeDeriverTests.java
@@ -0,0 +1,115 @@
+/*
+ * 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.planner.dag;
+
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.opensearch.analytics.planner.BasePlannerRulesTests;
+import org.opensearch.analytics.planner.rel.AnnotatedPredicate;
+import org.opensearch.analytics.planner.rel.OpenSearchFilter;
+import org.opensearch.analytics.spi.FilterTreeShape;
+
+import java.util.List;
+
+/**
+ * Unit tests for {@link FilterTreeShapeDeriver}.
+ */
+public class FilterTreeShapeDeriverTests extends BasePlannerRulesTests {
+
+ private static final String DRIVING = "datafusion";
+ private static final String ACCEPTING = "lucene";
+
+ public void testNoDelegation() {
+ // Single native predicate — no delegation
+ RexNode nativePred = annotated(DRIVING);
+ OpenSearchFilter filter = buildFilter(nativePred);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals("No delegation should return PLAIN", FilterTreeShape.NO_DELEGATION, shape);
+ }
+
+ public void testSingleDelegatedPredicate() {
+ // Single delegated predicate under implicit AND
+ RexNode delegated = annotated(ACCEPTING);
+ RexNode nativePred = annotated(DRIVING);
+ RexNode andNode = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, delegated);
+ OpenSearchFilter filter = buildFilter(andNode);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals(FilterTreeShape.CONJUNCTIVE, shape);
+ }
+
+ public void testMultipleDelegatedUnderAnd() {
+ // Multiple delegated predicates under AND — still SINGLE_AND
+ RexNode delegated1 = annotated(ACCEPTING);
+ RexNode delegated2 = annotated(ACCEPTING);
+ RexNode nativePred = annotated(DRIVING);
+ RexNode andNode = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, delegated1, delegated2);
+ OpenSearchFilter filter = buildFilter(andNode);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals(FilterTreeShape.CONJUNCTIVE, shape);
+ }
+
+ public void testOrWithDelegatedAndNative() {
+ // OR mixing delegated and native → MIXED_BOOLEAN
+ RexNode delegated = annotated(ACCEPTING);
+ RexNode nativePred = annotated(DRIVING);
+ RexNode orNode = rexBuilder.makeCall(SqlStdOperatorTable.OR, nativePred, delegated);
+ OpenSearchFilter filter = buildFilter(orNode);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals(FilterTreeShape.INTERLEAVED_BOOLEAN_EXPRESSION, shape);
+ }
+
+ public void testNotWithDelegated() {
+ // NOT wrapping delegated + native → MIXED_BOOLEAN
+ RexNode delegated = annotated(ACCEPTING);
+ RexNode nativePred = annotated(DRIVING);
+ RexNode andNode = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, delegated);
+ RexNode notNode = rexBuilder.makeCall(SqlStdOperatorTable.NOT, andNode);
+ OpenSearchFilter filter = buildFilter(notNode);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals(FilterTreeShape.INTERLEAVED_BOOLEAN_EXPRESSION, shape);
+ }
+
+ public void testOrWithOnlyDelegated() {
+ // OR with only delegated predicates (no driving backend) — SINGLE_AND (no mixing)
+ RexNode delegated1 = annotated(ACCEPTING);
+ RexNode delegated2 = annotated(ACCEPTING);
+ RexNode orNode = rexBuilder.makeCall(SqlStdOperatorTable.OR, delegated1, delegated2);
+ RexNode nativePred = annotated(DRIVING);
+ RexNode andNode = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, orNode);
+ OpenSearchFilter filter = buildFilter(andNode);
+
+ FilterTreeShape shape = FilterTreeShapeDeriver.derive(filter, DRIVING);
+ assertEquals(FilterTreeShape.CONJUNCTIVE, shape);
+ }
+
+ // ---- Helpers ----
+
+ private AnnotatedPredicate annotated(String backendId) {
+ RelDataType boolType = typeFactory.createJavaType(boolean.class);
+ RexNode literal = rexBuilder.makeLiteral(true);
+ return new AnnotatedPredicate(boolType, literal, List.of(backendId), 0);
+ }
+
+ private OpenSearchFilter buildFilter(RexNode condition) {
+ return new OpenSearchFilter(
+ cluster,
+ RelTraitSet.createEmpty(),
+ stubScan(mockTable("test_index", "col")),
+ condition,
+ List.of(DRIVING)
+ );
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java
index 1d2c4f81f16ba..1d1d6de31f427 100644
--- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java
@@ -38,7 +38,10 @@
import org.opensearch.analytics.spi.DelegatedPredicateSerializer;
import org.opensearch.analytics.spi.DelegationType;
import org.opensearch.analytics.spi.FieldStorageInfo;
+import org.opensearch.analytics.spi.FilterDelegationInstructionNode;
+import org.opensearch.analytics.spi.FilterTreeShape;
import org.opensearch.analytics.spi.FragmentConvertor;
+import org.opensearch.analytics.spi.InstructionType;
import org.opensearch.analytics.spi.ScalarFunction;
import java.nio.charset.StandardCharsets;
@@ -109,6 +112,10 @@ private void assertShardScanConverted(RecordingConvertor convertor, Stage stage)
assertEquals("test_index", convertor.shardScanTableName);
assertDoesntContainOperators(convertor.shardScanFragment, OPENSEARCH_OPERATORS);
assertDoesntContainOperators(convertor.shardScanFragment, ANNOTATION_MARKERS);
+ // Instruction assertions
+ StagePlan plan = stage.getPlanAlternatives().getFirst();
+ assertFalse("instructions must not be empty", plan.instructions().isEmpty());
+ assertEquals("first instruction must be SHARD_SCAN", InstructionType.SETUP_SHARD_SCAN, plan.instructions().getFirst().type());
}
private void assertReduceStageConverted(RecordingConvertor convertor, Stage stage) {
@@ -117,6 +124,14 @@ private void assertReduceStageConverted(RecordingConvertor convertor, Stage stag
assertTrue("convertFinalAggFragment must be called", convertor.finalAggCalled);
assertDoesntContainOperators(convertor.reduceFragment, OPENSEARCH_OPERATORS);
assertDoesntContainOperators(convertor.reduceFragment, ANNOTATION_MARKERS);
+ // Instruction assertions
+ StagePlan plan = stage.getPlanAlternatives().getFirst();
+ assertFalse("instructions must not be empty", plan.instructions().isEmpty());
+ assertEquals(
+ "reduce stage must have FINAL_AGGREGATE",
+ InstructionType.SETUP_FINAL_AGGREGATE,
+ plan.instructions().getFirst().type()
+ );
}
// ---- Single-stage query shapes ----
@@ -350,9 +365,10 @@ private void assertDelegationResult(
int expectedDelegatedCount,
boolean expectPlaceholder,
boolean expectNativeEquals,
- List expectedFunctions
+ List expectedFunctions,
+ FilterTreeShape expectedTreeShape
) {
- assertEquals("delegatedQueries count", expectedDelegatedCount, plan.delegatedQueries().size());
+ assertEquals("delegatedQueries count", expectedDelegatedCount, plan.delegatedExpressions().size());
assertEquals("serializer call count", expectedDelegatedCount, serializer.callCount);
assertEquals("serialized functions", expectedFunctions, serializer.serializedFunctions);
@@ -379,6 +395,26 @@ private void assertDelegationResult(
// No annotation markers should survive stripping
assertDoesntContainOperators(dfConvertor.shardScanFragment, ANNOTATION_MARKERS);
+
+ // Instruction assertions: delegation plans must have SHARD_SCAN + FILTER_DELEGATION_FOR_INDEX
+ if (expectedDelegatedCount > 0) {
+ assertTrue(
+ "delegation plan must have FILTER_DELEGATION_FOR_INDEX instruction",
+ plan.instructions().stream().anyMatch(node -> node.type() == InstructionType.SETUP_FILTER_DELEGATION_FOR_INDEX)
+ );
+ FilterDelegationInstructionNode filterInstruction = (FilterDelegationInstructionNode) plan.instructions()
+ .stream()
+ .filter(node -> node.type() == InstructionType.SETUP_FILTER_DELEGATION_FOR_INDEX)
+ .findFirst()
+ .orElseThrow();
+ assertEquals("delegatedPredicateCount in instruction", expectedDelegatedCount, filterInstruction.getDelegatedPredicateCount());
+ assertEquals(
+ "delegatedExpressions in instruction must match plan",
+ plan.delegatedExpressions().size(),
+ filterInstruction.getDelegatedQueries().size()
+ );
+ assertEquals("treeShape in instruction", expectedTreeShape, filterInstruction.getTreeShape());
+ }
}
// ---- Single predicate ----
@@ -389,7 +425,7 @@ public void testSingleDelegatedPredicate() {
RecordingSerializer serializer = new RecordingSerializer();
QueryDAG dag = buildSingleFieldDelegationDag(makeFullTextCall(MATCH_PHRASE_FUNCTION, 0, "hello world"), dfConvertor, serializer);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 1, true, false, List.of("MATCH_PHRASE"));
+ assertDelegationResult(plan, dfConvertor, serializer, 1, true, false, List.of("MATCH_PHRASE"), FilterTreeShape.CONJUNCTIVE);
}
/** Single native equals — no delegation, empty delegatedQueries. */
@@ -398,7 +434,7 @@ public void testSingleNativePredicate() {
RecordingSerializer serializer = new RecordingSerializer();
QueryDAG dag = buildTwoFieldDelegationDag(makeEquals(0, SqlTypeName.INTEGER, 200), dfConvertor, serializer);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 0, false, true, List.of());
+ assertDelegationResult(plan, dfConvertor, serializer, 0, false, true, List.of(), FilterTreeShape.NO_DELEGATION);
}
// ---- AND conditions ----
@@ -413,7 +449,7 @@ public void testAndNativeAndDelegated() {
serializer
);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 1, true, true, List.of("MATCH_PHRASE"));
+ assertDelegationResult(plan, dfConvertor, serializer, 1, true, true, List.of("MATCH_PHRASE"), FilterTreeShape.CONJUNCTIVE);
}
/** AND(delegated, delegated) — both replaced, two entries in delegatedQueries. */
@@ -426,7 +462,16 @@ public void testAndTwoDelegated() {
serializer
);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 2, true, false, List.of("MATCH_PHRASE", "FUZZY"));
+ assertDelegationResult(
+ plan,
+ dfConvertor,
+ serializer,
+ 2,
+ true,
+ false,
+ List.of("MATCH_PHRASE", "FUZZY"),
+ FilterTreeShape.CONJUNCTIVE
+ );
}
// ---- OR conditions ----
@@ -445,7 +490,16 @@ public void testOrNativeAndDelegated() {
serializer
);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 1, true, true, List.of("MATCH_PHRASE"));
+ assertDelegationResult(
+ plan,
+ dfConvertor,
+ serializer,
+ 1,
+ true,
+ true,
+ List.of("MATCH_PHRASE"),
+ FilterTreeShape.INTERLEAVED_BOOLEAN_EXPRESSION
+ );
assertTrue("OR structure should be preserved", RelOptUtil.toString(dfConvertor.shardScanFragment).contains("OR"));
}
@@ -467,7 +521,7 @@ public void testInterleavedAndOrNot() {
RexNode condition = makeAnd(makeEquals(0, SqlTypeName.INTEGER, 200), orClause);
QueryDAG dag = buildTwoFieldDelegationDag(condition, dfConvertor, serializer);
StagePlan plan = leafStage(dag).getPlanAlternatives().getFirst();
- assertDelegationResult(plan, dfConvertor, serializer, 2, true, true, List.of("MATCH_PHRASE", "FUZZY"));
+ assertDelegationResult(plan, dfConvertor, serializer, 2, true, true, List.of("MATCH_PHRASE", "FUZZY"), FilterTreeShape.CONJUNCTIVE);
String strippedPlan = RelOptUtil.toString(dfConvertor.shardScanFragment);
assertTrue("AND structure should be preserved", strippedPlan.contains("AND"));
assertTrue("OR structure should be preserved", strippedPlan.contains("OR"));