Stream Arrow batches on the data-node fragment execution path

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

@@ -8,22 +8,34 @@
 
 package org.opensearch.analytics.backend;
 
+import org.apache.arrow.vector.VectorSchemaRoot;
+
 import java.util.List;
 
 /**
- * Read-only view of a single record batch. Provides field names, row count,
- * and positional access to field values.
+ * Read-only view of a single record batch.
  * <p>
  * A batch is only valid until the next call to {@link java.util.Iterator#next()}
  * on the parent stream's iterator. The underlying data buffers may be reused
  * across batches, so callers must extract all needed values before advancing
  * the iterator. Accessing a batch after the iterator has advanced may throw
  * {@link IllegalStateException}.
+ * <p>
+ * Primary shape is the Arrow {@link VectorSchemaRoot} returned by
+ * {@link #getArrowRoot()} — the native columnar representation used by the
+ * streaming transport (zero-copy over gRPC). Row-oriented accessors
+ * ({@link #getFieldNames()}, {@link #getRowCount()}, {@link #getFieldValue})
+ * are a convenience view over the same data.
  *
  * @opensearch.internal
  */
 public interface EngineResultBatch {
 
+    /**
+     * The Arrow VSR backing this batch
+     */
+    VectorSchemaRoot getArrowRoot();
+
     /**
      * Ordered list of field (column) names in this batch.
      */

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

@@ -8,8 +8,9 @@
 
 package org.opensearch.analytics.backend;
 
-import org.opensearch.action.search.SearchShardTask;
+import org.apache.arrow.memory.BufferAllocator;
 import org.opensearch.index.engine.exec.IndexReaderProvider.Reader;
+import org.opensearch.tasks.Task;
 
 /**
  * Execution context carrying reader and plan state through
@@ -21,23 +22,24 @@ public class ExecutionContext {
 
     private final String tableName;
     private final Reader reader;
-    private final SearchShardTask task;
+    private final Task task;
     private byte[] fragmentBytes;
+    private BufferAllocator allocator;
 
     /**
      * Constructs an execution context.
      * @param tableName the target table name
-     * @param task the search shard task
+     * @param task the transport-created task for this fragment execution
      * @param reader the data-format aware reader
      */
-    public ExecutionContext(String tableName, SearchShardTask task, Reader reader) {
+    public ExecutionContext(String tableName, Task task, Reader reader) {
         this.tableName = tableName;
         this.task = task;
         this.reader = reader;
     }
 
-    /** Returns the search shard task. */
-    public SearchShardTask getTask() {
+    /** Returns the transport-created task for this fragment execution. */
+    public Task getTask() {
         return task;
     }
 
@@ -60,4 +62,14 @@ public byte[] getFragmentBytes() {
     public void setFragmentBytes(byte[] fragmentBytes) {
         this.fragmentBytes = fragmentBytes;
     }
+
+    /** Returns the caller-provided allocator for producing Arrow result buffers. */
+    public BufferAllocator getAllocator() {
+        return allocator;
+    }
+
+    /** Sets the caller-provided allocator. The caller owns its lifecycle; the engine must not close it. */
+    public void setAllocator(BufferAllocator allocator) {
+        this.allocator = allocator;
+    }
 }

sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/StreamingCoordinatorReduceIT.java

@@ -0,0 +1,145 @@
+/*
+ * 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.Version;
+import org.opensearch.action.admin.indices.create.CreateIndexResponse;
+import org.opensearch.analytics.AnalyticsPlugin;
+import org.opensearch.arrow.flight.transport.FlightStreamPlugin;
+import org.opensearch.be.lucene.LucenePlugin;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
+import org.opensearch.composite.CompositeDataFormatPlugin;
+import org.opensearch.parquet.ParquetDataFormatPlugin;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.plugins.PluginInfo;
+import org.opensearch.ppl.TestPPLPlugin;
+import org.opensearch.ppl.action.PPLRequest;
+import org.opensearch.ppl.action.PPLResponse;
+import org.opensearch.ppl.action.UnifiedPPLExecuteAction;
+import org.opensearch.test.OpenSearchIntegTestCase;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+import static org.opensearch.common.util.FeatureFlags.STREAM_TRANSPORT;
+
+/**
+ * Streaming variant of {@link CoordinatorReduceIT}: same 2-shard parquet-backed index and deterministic
+ * dataset, but with Arrow Flight RPC enabled via {@link FeatureFlags#STREAM_TRANSPORT}. Exercises the
+ * shard-fragment → Flight → {@code DatafusionReduceSink.feed} handoff that previously failed with
+ * {@code "A buffer can only be associated between two allocators that share the same root"} on
+ * multi-shard queries.
+ */
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 2)
+public class StreamingCoordinatorReduceIT extends OpenSearchIntegTestCase {
+
+    private static final String INDEX = "coord_reduce_streaming_e2e";
+    private static final int NUM_SHARDS = 2;
+    private static final int DOCS_PER_SHARD = 10;
+    /** Constant `value` for every doc — deterministic assertion independent of shard routing. */
+    private static final int VALUE = 7;
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(TestPPLPlugin.class, FlightStreamPlugin.class, CompositeDataFormatPlugin.class, LucenePlugin.class);
+    }
+
+    @Override
+    protected Collection<PluginInfo> additionalNodePlugins() {
+        return List.of(
+            classpathPlugin(AnalyticsPlugin.class, Collections.emptyList()),
+            classpathPlugin(ParquetDataFormatPlugin.class, Collections.emptyList()),
+            classpathPlugin(DataFusionPlugin.class, List.of(AnalyticsPlugin.class.getName()))
+        );
+    }
+
+    private static PluginInfo classpathPlugin(Class<? extends Plugin> pluginClass, List<String> extendedPlugins) {
+        return new PluginInfo(
+            pluginClass.getName(),
+            "classpath plugin",
+            "NA",
+            Version.CURRENT,
+            "1.8",
+            pluginClass.getName(),
+            null,
+            extendedPlugins,
+            false
+        );
+    }
+
+    @Override
+    protected Settings nodeSettings(int nodeOrdinal) {
+        return Settings.builder()
+            .put(super.nodeSettings(nodeOrdinal))
+            .put(FeatureFlags.PLUGGABLE_DATAFORMAT_EXPERIMENTAL_FLAG, true)
+            .build();
+    }
+
+    /**
+     * {@code source = T} on a 2-shard parquet-backed index with streaming enabled exercises the
+     * coordinator reduce sink's cross-plugin VectorSchemaRoot handoff. Before the allocator-root
+     * unification fix, this failed with an Arrow {@code associate} mismatch.
+     */
+    @LockFeatureFlag(STREAM_TRANSPORT)
+    public void testBaselineScanAcrossShards() throws Exception {
+        createParquetBackedIndex();
+        indexDeterministicDocs();
+
+        PPLResponse response = executePPL("source = " + INDEX);
+
+        assertNotNull("PPLResponse must not be null", response);
+        assertTrue("columns must contain 'value', got " + response.getColumns(), response.getColumns().contains("value"));
+
+        int expectedRows = NUM_SHARDS * DOCS_PER_SHARD;
+        assertEquals("all docs across shards must be returned", expectedRows, response.getRows().size());
+
+        int idx = response.getColumns().indexOf("value");
+        for (Object[] row : response.getRows()) {
+            Object cell = row[idx];
+            assertNotNull("value cell must not be null", cell);
+            assertEquals("every doc has value=" + VALUE, (long) VALUE, ((Number) cell).longValue());
+        }
+    }
+
+    private void createParquetBackedIndex() {
+        Settings indexSettings = Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, NUM_SHARDS)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            .put("index.pluggable.dataformat.enabled", true)
+            .put("index.pluggable.dataformat", "composite")
+            .put("index.composite.primary_data_format", "parquet")
+            .putList("index.composite.secondary_data_formats")
+            .build();
+
+        CreateIndexResponse response = client().admin()
+            .indices()
+            .prepareCreate(INDEX)
+            .setSettings(indexSettings)
+            .setMapping("value", "type=integer")
+            .get();
+        assertTrue("index creation must be acknowledged", response.isAcknowledged());
+        ensureGreen(INDEX);
+    }
+
+    private void indexDeterministicDocs() {
+        int total = NUM_SHARDS * DOCS_PER_SHARD;
+        for (int i = 0; i < total; i++) {
+            client().prepareIndex(INDEX).setId(String.valueOf(i)).setSource("value", VALUE).get();
+        }
+        client().admin().indices().prepareRefresh(INDEX).get();
+        client().admin().indices().prepareFlush(INDEX).get();
+    }
+
+    private PPLResponse executePPL(String ppl) {
+        return client().execute(UnifiedPPLExecuteAction.INSTANCE, new PPLRequest(ppl)).actionGet();
+    }
+}

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

@@ -175,7 +175,7 @@ public SearchExecEngineProvider getSearchExecEngineProvider() {
                 throw new IllegalStateException("No DatafusionReader available in the acquired reader");
             }
             DatafusionContext context = new DatafusionContext(ctx.getTask(), dfReader, dataFusionService.getNativeRuntime());
-            DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context, dataFusionService::newChildAllocator);
+            DatafusionSearchExecEngine engine = new DatafusionSearchExecEngine(context);
             engine.prepare(ctx);
             return engine;
         };

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

@@ -8,10 +8,10 @@
 
 package org.opensearch.be.datafusion;
 
-import org.opensearch.action.search.SearchShardTask;
 import org.opensearch.be.datafusion.nativelib.StreamHandle;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.search.SearchExecutionContext;
+import org.opensearch.tasks.Task;
 
 import java.io.IOException;
 
@@ -30,15 +30,15 @@ public class DatafusionContext implements SearchExecutionContext<DatafusionSearc
     private final NativeRuntimeHandle nativeRuntime;
     private DatafusionQuery datafusionQuery;
     private StreamHandle streamHandle;
-    private SearchShardTask task;
+    private Task task;
 
     /**
      * Creates a DataFusion execution context
      * @param task the search shard task
      * @param reader the DataFusion reader providing index data
      * @param nativeRuntime handle to the native DataFusion runtime
      */
-    public DatafusionContext(SearchShardTask task, DatafusionReader reader, NativeRuntimeHandle nativeRuntime) {
+    public DatafusionContext(Task task, DatafusionReader reader, NativeRuntimeHandle nativeRuntime) {
         this.task = task;
         this.engineSearcher = new DatafusionSearcher(reader.getReaderHandle());
         this.nativeRuntime = nativeRuntime;
@@ -101,7 +101,7 @@ public void setStreamHandle(StreamHandle streamHandle) {
     }
 
     @Override
-    public SearchShardTask task() {
+    public Task task() {
         return task;
     }