opensearch-project · mch2 · May 28, 2026 · May 25, 2026 · May 25, 2026 · May 25, 2026
@@ -8,6 +8,11 @@
 
 package org.opensearch.analytics.spi;
 
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.opensearch.analytics.backend.EngineResultStream;
+import org.opensearch.index.engine.exec.IndexReaderProvider.Reader;
+
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -134,6 +139,20 @@ default Map<Long, QueryExecutionMetrics> getTopQueriesByMemory() {
         return Collections.emptyMap();
     }
 
+    /**
+     * QTF fetch phase: reads specific rows by global row ID.
+     * Row IDs are passed as a BigIntVector for zero-copy transfer to native.
+     *
+     * @param reader the index reader for the target shard
+     * @param rowIdVector Arrow BigIntVector containing global row IDs
+     * @param columns column names to read
+     * @param allocator Arrow buffer allocator for result import
+     * @return a result stream containing the requested rows
+     */
+    default EngineResultStream fetchByRowIds(Reader reader, BigIntVector rowIdVector, String[] columns, BufferAllocator allocator) {
+        throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]");
+    }
+
     /**
      * Install a thread tracker for attribution of delegation callbacks executing on foreign threads.
      * Called after {@link #configureFilterDelegation}. Pass {@code null} to clear.

@@ -32,6 +32,24 @@ public interface ExchangeSink {
      */
     void feed(VectorSchemaRoot batch);
 
+    /**
+     * Ingest an Arrow batch with a per-source ordinal — the index of the
+     * producer task within its stage's resolved target list (e.g.
+     * {@link org.opensearch.analytics.spi.ExchangeSink} consumed via
+     * {@code ShardExecutionTarget.ordinal()}).
+     *
+     * <p>Default implementation drops the ordinal and falls through to
+     * {@link #feed(VectorSchemaRoot)}. Sinks that need to discriminate
+     * batches by producer (e.g. Late Materialization, where the ordinal is
+     * stamped onto each batch as a column) override this method.
+     *
+     * <p>Producers that have a meaningful per-task ordinal call this overload;
+     * producers without one continue to call {@link #feed(VectorSchemaRoot)}.
+     */
+    default void feed(VectorSchemaRoot batch, int sourceOrdinal) {
+        feed(batch);
+    }
+
     /**
      * Signal that no more batches will be fed. Releases resources.
      */

@@ -10,6 +10,7 @@
 
 import org.apache.calcite.sql.type.SqlTypeName;
 
+import java.util.LinkedHashSet;
 import java.util.List;
 
 /**
@@ -28,6 +29,7 @@ public class FieldStorageInfo {
     private final List<String> indexFormats;
     private final List<String> storedFieldFormats;
     private final boolean derived;
+    private final LinkedHashSet<String> dependsOnPhysicalCols;
 
     public FieldStorageInfo(
         String fieldName,
@@ -37,6 +39,21 @@ public FieldStorageInfo(
         List<String> indexFormats,
         List<String> storedFieldFormats,
         boolean derived
+    ) {
+        // Default: no physical-col dependencies. Physical fields aren't "derived from"
+        // anything; derived fields' deps are supplied by the caller via the 8-arg ctor.
+        this(fieldName, mappingType, fieldType, docValueFormats, indexFormats, storedFieldFormats, derived, new LinkedHashSet<>());
+    }
+
+    public FieldStorageInfo(
+        String fieldName,
+        String mappingType,
+        FieldType fieldType,
+        List<String> docValueFormats,
+        List<String> indexFormats,
+        List<String> storedFieldFormats,
+        boolean derived,
+        LinkedHashSet<String> dependsOnPhysicalCols
     ) {
         this.fieldName = fieldName;
         this.mappingType = mappingType;
@@ -45,19 +62,30 @@ public FieldStorageInfo(
         this.indexFormats = indexFormats;
         this.storedFieldFormats = storedFieldFormats;
         this.derived = derived;
+        this.dependsOnPhysicalCols = dependsOnPhysicalCols;
     }
 
-    /** Creates a derived column (agg result, expression) with no physical storage.
-     *  FieldType inferred from SqlTypeName. */
+    /** Creates a derived column (agg result, expression) with no physical storage and no deps.
+     *  FieldType inferred from SqlTypeName. Use {@link #derivedColumn(String, SqlTypeName, LinkedHashSet)}
+     *  when the caller can supply the underlying physical-column dependencies. */
     public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName) {
+        return derivedColumn(fieldName, sqlTypeName, new LinkedHashSet<>());
+    }
+
+    /** Creates a derived column with explicit physical-column dependencies — the
+     *  TableScan-level fields whose values flow into this column's computation, in
+     *  first-appearance order. {@link LinkedHashSet} makes both the ordering invariant
+     *  and the no-duplicates invariant explicit at the type level. */
+    public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName, LinkedHashSet<String> dependsOnPhysicalCols) {
         return new FieldStorageInfo(
             fieldName,
             sqlTypeName.getName(),
             FieldType.fromSqlTypeName(sqlTypeName),
             List.of(),
             List.of(),
             List.of(),
-            true
+            true,
+            dependsOnPhysicalCols
         );
     }
 
@@ -88,6 +116,29 @@ public boolean isDerived() {
         return derived;
     }
 
+    /**
+     * Names of the TableScan-level (physical) columns whose values flow into this column's
+     * computation, in first-appearance order. Empty for physical columns (they aren't
+     * "derived from" anything — their identity is their own field name) and for derived
+     * columns with no physical inputs (e.g. {@code COUNT(*)}, pure-literal projects).
+     * For other derived columns it is the union of underlying physical cols across the
+     * expression / agg-call tree.
+     *
+     * <p>Used by the QTF rewriter to derive the fetch list off the topmost operator's FSI
+     * without re-walking RexNodes from scratch. {@link LinkedHashSet} makes both the
+     * ordering invariant and the no-duplicates invariant explicit at the type level.
+     *
+     * <p>TODO: today we use string field names because they stay stable across narrowed-scan
+     * rewrites and (future) Join/Union plans where int ordinals across multiple TableScans
+     * become ambiguous. For very large plans the per-FSI string set can become a memory
+     * hotspot — consider switching to int ordinals (rooted to the originating TableScan's
+     * rowType) when single-scan throughput dominates and stability across rewrites can be
+     * traded off.
+     */
+    public LinkedHashSet<String> getDependsOnPhysicalCols() {
+        return dependsOnPhysicalCols;
+    }
+
     /**
      * Resolves a field by index from a fieldStorageInfos list, validating bounds and field type.
      * Throws if the index is out of bounds or the field type is unrecognized.

@@ -9,6 +9,7 @@
 package org.opensearch.analytics.spi;
 
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.type.RelDataType;
 
 /**
  * Fragment conversion API for backend plugins.
@@ -81,4 +82,19 @@ default byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByt
     default byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) {
         throw new UnsupportedOperationException("attachFragmentOnTop not implemented for this backend");
     }
+
+    /**
+     * Builds a schema-only stub plan describing a stage's output partition: a single
+     * {@code Read { named_table: "input-<childStageId>"; base_schema: rowType }}, no
+     * operators above. Used for stages whose runtime is non-Substrait (e.g. QTF
+     * scatter-gather) but whose parent reduce sink still needs a partition schema for
+     * {@code registerPartitionStream}.
+     *
+     * @param childStageId stage id whose output schema this stub describes
+     * @param rowType      the partition's row type
+     * @return serialized plan bytes
+     */
+    default byte[] convertSchemaOnlyRead(int childStageId, RelDataType rowType) {
+        throw new UnsupportedOperationException("convertSchemaOnlyRead not implemented for this backend");
+    }
 }
@@ -25,8 +25,11 @@ public interface FragmentInstructionHandlerFactory {
 
     // ── Coordinator-side: create instruction nodes ──
 
-    /** Creates a shard scan instruction node. */
-    Optional<InstructionNode> createShardScanNode();
+    /**
+     * Creates a shard scan instruction node. {@code requestsRowIds} signals that the scan
+     * must emit shard-global {@code __row_id__} values (QTF query phase).
+     */
+    Optional<InstructionNode> createShardScanNode(boolean requestsRowIds);
 
     /** Creates a filter delegation instruction node with the given delegation metadata. */
     Optional<InstructionNode> createFilterDelegationNode(
@@ -35,8 +38,17 @@ Optional<InstructionNode> createFilterDelegationNode(
         List<DelegatedExpression> delegatedQueries
     );
 
-    /** Creates a shard scan with delegation instruction node — combines scan setup with delegation config. */
-    Optional<InstructionNode> createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount);
+    /**
+     * Creates a shard scan with delegation instruction node — combines scan setup with
+     * delegation config. {@code requestsRowIds} signals that the scan must emit shard-global
+     * {@code __row_id__} values (QTF query phase). Backends that don't support QTF should
+     * return {@link Optional#empty()} when {@code requestsRowIds} is true.
+     */
+    Optional<InstructionNode> createShardScanWithDelegationNode(
+        FilterTreeShape treeShape,
+        int delegatedPredicateCount,
+        boolean requestsRowIds
+    );
 
     /** Creates a partial aggregate instruction node. */
     Optional<InstructionNode> createPartialAggregateNode();

@@ -10,21 +10,37 @@
 
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
 
 import java.io.IOException;
 
 /**
  * Instruction node for base shard scan setup — reader acquisition, SessionContext creation,
- * default table provider registration.
+ * table provider registration. {@code requestsRowIds} signals that the shard scan needs to
+ * emit shard-global {@code __row_id__} values (QTF query phase). Inherited by
+ * {@link ShardScanWithDelegationInstructionNode} so the same flag applies whether or not
+ * filter delegation is in play — QTF and delegation are orthogonal concerns.
  *
  * @opensearch.internal
  */
-public class ShardScanInstructionNode implements InstructionNode {
+public class ShardScanInstructionNode implements InstructionNode, Writeable {
 
-    public ShardScanInstructionNode() {}
+    private final boolean requestsRowIds;
+
+    public ShardScanInstructionNode() {
+        this(false);
+    }
+
+    public ShardScanInstructionNode(boolean requestsRowIds) {
+        this.requestsRowIds = requestsRowIds;
+    }
 
     public ShardScanInstructionNode(StreamInput in) throws IOException {
-        // No fields to read
+        this.requestsRowIds = in.readBoolean();
+    }
+
+    public boolean requestsRowIds() {
+        return requestsRowIds;
     }
 
     @Override
@@ -34,6 +50,6 @@ public InstructionType type() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        // No fields to write
+        out.writeBoolean(requestsRowIds);
     }
 }
@@ -27,6 +27,11 @@ public class ShardScanWithDelegationInstructionNode extends ShardScanInstruction
     private final int delegatedPredicateCount;
 
     public ShardScanWithDelegationInstructionNode(FilterTreeShape treeShape, int delegatedPredicateCount) {
+        this(treeShape, delegatedPredicateCount, false);
+    }
+
+    public ShardScanWithDelegationInstructionNode(FilterTreeShape treeShape, int delegatedPredicateCount, boolean requestsRowIds) {
+        super(requestsRowIds);
         this.treeShape = treeShape;
         this.delegatedPredicateCount = delegatedPredicateCount;
     }

@@ -107,3 +107,7 @@ harness = false
 [[bench]]
 name = "cross_rt_throughput_bench"
 harness = false
+
+[[bench]]
+name = "row_id_bench"
+harness = false
@@ -5,7 +5,7 @@ use datafusion::execution::memory_pool::GreedyMemoryPool;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use futures::TryStreamExt;
 use object_store::local::LocalFileSystem;
-use object_store::ObjectStore;
+use object_store::{ObjectStore, ObjectStoreExt};
 use opensearch_datafusion::api::DataFusionRuntime;
 use opensearch_datafusion::query_executor;
 use opensearch_datafusion::runtime_manager::RuntimeManager;
@@ -104,7 +104,16 @@ fn bench_execute_query(c: &mut Criterion) {
                     let exec = mgr.cpu_executor();
                     async {
                         let ptr = query_executor::execute_query(
-                            url, metas, "t".into(), plan, &df_runtime, exec, None, &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(),
+                            url,
+                            metas,
+                            "t".into(),
+                            plan,
+                            &df_runtime,
+                            exec,
+                            None,
+                            &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(),
+                            0,
+                            Arc::new(LocalFileSystem::new()) as Arc<dyn ObjectStore>,
                         ).await.unwrap();
                         // Consume and free the stream
                         let mut stream = unsafe {
@@ -151,6 +160,8 @@ fn bench_stream_next(c: &mut Criterion) {
                     None,
                     &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(
                     ),
+                    0,
+                    Arc::new(LocalFileSystem::new()) as Arc<dyn ObjectStore>,
                 )
                 .await
                 .unwrap();
@@ -199,6 +210,8 @@ fn bench_aggregation(c: &mut Criterion) {
                     None,
                     &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(
                     ),
+                    0,
+                    Arc::new(LocalFileSystem::new()) as Arc<dyn ObjectStore>,
                 )
                 .await
                 .unwrap();