opensearch-project · sandeshkr419 · Jun 13, 2026 · Jun 9, 2026 · Jun 9, 2026 · Jun 10, 2026
@@ -34,6 +34,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef;
 import org.opensearch.analytics.planner.rules.ExtractLiteralAggRule;
+import org.opensearch.analytics.planner.rules.OpenSearchAggLiteralArgProjectSplitRule;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateReduceRule;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateRule;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateSplitRule;
@@ -53,6 +54,8 @@
 import org.opensearch.analytics.planner.rules.OpenSearchUnionSplitRule;
 import org.opensearch.analytics.planner.rules.OpenSearchValuesRule;
 
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.List;
 import java.util.Optional;
 
@@ -101,6 +104,7 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con
         modifiedRelNode = decomposeAggregates(modifiedRelNode, listener);
         modifiedRelNode = mark(modifiedRelNode, context, listener);
         LOGGER.debug("After marking:\n{}", RelOptUtil.toString(modifiedRelNode));
+        modifiedRelNode = splitAggLiteralArgProject(modifiedRelNode, listener);
         // TODO(combine-delegated-predicates): a post-marking HEP rule should fuse same-backend
         // AND-sibling AnnotatedPredicates into one combined predicate per group, collapsing N
         // FFM round-trips per RG into one. Blocked on two open design points:
@@ -248,6 +252,20 @@ private static RelNode extractLiteralAgg(RelNode input, RuleProfilingListener li
         return HepPhase.named("literal-agg-extract").addRuleInstance(new ExtractLiteralAggRule()).run(input, listener);
     }
 
+    /**
+     * Phase 1c': duplicate an aggregate's literal-config-arg Project (e.g. percentile's {@code 50})
+     * into a pinned upper copy (literal stays with the aggregate) over an unpinned physical-only
+     * lower copy (pushes below the ExchangeReducer). Runs AFTER marking — operates on
+     * {@code OpenSearch*} nodes and emits a pinned {@code OpenSearchProject} whose
+     * {@code computeSelfCost} forces the CBO-inserted ER below it, keeping the literal in the
+     * coordinator fragment for the DataFusion substrait converter. Placed after marking so the
+     * pre-marking {@code PROJECT_MERGE} cannot re-fuse the two copies. See
+     * {@link OpenSearchAggLiteralArgProjectSplitRule}.
+     */
+    private static RelNode splitAggLiteralArgProject(RelNode input, RuleProfilingListener listener) {
+        return HepPhase.named("agg-literal-arg-split").addRuleInstance(new OpenSearchAggLiteralArgProjectSplitRule()).run(input, listener);
+    }
+
     /**
      * Phase 1a: constant-expression reduction on Filter and Project predicates. Kept in
      * its own phase so {@code ProjectReduceExpressionsRule} cannot use a downstream
@@ -277,15 +295,16 @@ private static RelNode reduceExpressions(RelNode input, RuleProfilingListener li
     private static RelNode pushdownRules(RelNode input, RuleProfilingListener listener) {
         return HepPhase.named("pushdown-rules")
             .bottomUp()
-            // Transposes (filter-into-* and sort-into-project) cascade together within
-            // one fixpoint, alongside PROJECT_MERGE which collapses the intermediate
-            // adjacent Projects that SORT_PROJECT_TRANSPOSE produces. FILTER_MERGE
-            // runs as its own instruction so it only fires after the transposes have
-            // settled — that way any auto-injected NOT NULL collapses with the user's
-            // WHERE on the post-pushdown filter, not on a half-pushed intermediate.
-            // SORT_PROJECT_TRANSPOSE + PROJECT_MERGE feed the QTF (late-materialization)
-            // rewriter by lifting Project above Sort so it sees a single Project layer
-            // above the anchor.
+            // Transposes (filter-into-*) cascade together within one fixpoint, alongside
+            // PROJECT_MERGE which collapses adjacent Projects. FILTER_MERGE runs as its own
+            // instruction so it only fires after the transposes have settled — that way any
+            // auto-injected NOT NULL collapses with the user's WHERE on the post-pushdown filter,
+            // not on a half-pushed intermediate.
+            //
+            // SORT_PROJECT_TRANSPOSE is intentionally omitted: lifting Project above Sort puts it
+            // above the Exchange, defeating projection pushdown. Keeping it below lets DataFusion
+            // prune the scan. QTF relocates the below-Sort Project above its wrapper itself.
+            //
             // SORT_REMOVE_REDUNDANT drops a Sort/LIMIT whose input is provably bounded to
             // within the limit (e.g. a collation-less `head N` or a sort over a scalar
             // aggregate, getMaxRowCount <= 1): a no-op that the marking rule must not have to
@@ -297,7 +316,6 @@ private static RelNode pushdownRules(RelNode input, RuleProfilingListener listen
                     CoreRules.FILTER_PROJECT_TRANSPOSE,
                     CoreRules.FILTER_AGGREGATE_TRANSPOSE,
                     CoreRules.FILTER_INTO_JOIN,
-                    CoreRules.SORT_PROJECT_TRANSPOSE,
                     CoreRules.PROJECT_MERGE,
                     CoreRules.LIMIT_MERGE,
                     CoreRules.SORT_REMOVE_REDUNDANT
@@ -390,7 +408,13 @@ private static RelNode cbo(RelNode marked, RelNode rawRelNode, PlannerContext co
             if (!copied.getTraitSet().equals(desiredTraits)) {
                 volcanoPlanner.setRoot(volcanoPlanner.changeTraits(copied, desiredTraits));
             }
-            return volcanoPlanner.findBestExp();
+            RelNode best = volcanoPlanner.findBestExp();
+            if (LOGGER.isDebugEnabled()) {
+                StringWriter sw = new StringWriter();
+                volcanoPlanner.dump(new PrintWriter(sw));
+                LOGGER.debug("Volcano memo:\n{}", sw);
+            }
+            return best;
         } finally {
             if (listener != null) listener.endPhase("cbo");
         }

@@ -283,16 +283,15 @@ private static RelNode cutAtExchange(
             new Stage(childStageId, childFragment, grandchildren, reducer.getExchangeInfo(), childSinkProvider, targetResolver)
         );
 
-        // Use the reducer's OUTPUT rowType so QTF's appended ___ugsi (set on erRowType by the
-        // rewriter) flows into the parent stage's partition schema. No-op for non-QTF reducers.
-        OpenSearchRelNode reducerInput = (OpenSearchRelNode) reducer.getInput();
+        // Source both rowType and FSI from the reducer so they stay aligned 1:1 (its input lacks
+        // QTF's ___ugsi entry). No-op for non-QTF reducers.
         OpenSearchStageInputScan stageInput = new OpenSearchStageInputScan(
             reducer.getCluster(),
             reducer.getTraitSet(),
             childStageId,
             reducer.getRowType(),
             reducer.getViableBackends(),
-            reducerInput.getOutputFieldStorage()
+            reducer.getOutputFieldStorage()
         );
         return new OpenSearchExchangeReducer(
             reducer.getCluster(),

@@ -261,10 +261,6 @@ private static boolean containsEngineNativeAggregate(RelNode root, AggregateMode
         return false;
     }
 
-    private static boolean isPureReorderProject(org.apache.calcite.rel.core.Project project) {
-        return project.getProjects().stream().allMatch(e -> e instanceof org.apache.calcite.rex.RexInputRef);
-    }
-
     /**
      * Accumulates serialized delegated query bytes during fragment conversion.
      *
@@ -534,14 +530,10 @@ private static byte[] convertReduceNode(
             }
 
             // Single-input operator above the final-fragment boundary — convert child first, then attach.
+            // A pure-reorder Project above an engine-native-merge FINAL is emitted like any other operator;
+            // dropping it would strand operators above it (e.g. Sort) with the post-reorder schema over
+            // un-reordered data, corrupting the final column order.
             byte[] innerBytes = convertReduceNode(node.getInputs().getFirst(), convertor, false, delegationBytes);
-            // Skip pure-reorder Project above engine-native-merge FINAL — DataFusion's substrait
-            // consumer can't bind reordered field names against the FINAL aggregate's state columns.
-            if (node instanceof org.apache.calcite.rel.core.Project p
-                && isPureReorderProject(p)
-                && containsEngineNativeAggregate(node.getInputs().getFirst(), AggregateMode.FINAL)) {
-                return innerBytes;
-            }
             return convertor.attachFragmentOnTop(strippedNode, innerBytes);
         }
         throw new IllegalStateException("Unexpected reduce stage node: " + node.getClass().getSimpleName());

@@ -43,6 +43,13 @@ public class OpenSearchExchangeReducer extends ConverterImpl implements OpenSear
      */
     private final RelDataType overrideRowType;
 
+    /**
+     * Field storage matching {@link #overrideRowType} 1:1, supplied by QTF alongside the override
+     * so the {@code ___ugsi} column it declares has a corresponding {@link FieldStorageInfo}. Null
+     * whenever {@code overrideRowType} is null; non-null only on QTF-rebuilt ERs.
+     */
+    private final List<FieldStorageInfo> overrideStorage;
+
     /** Convenience constructor — defaults to {@link ExchangeInfo#singleton()}. */
     public OpenSearchExchangeReducer(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, List<String> viableBackends) {
         this(cluster, traitSet, input, viableBackends, ExchangeInfo.singleton(), null);
@@ -71,6 +78,23 @@ public OpenSearchExchangeReducer(
         List<String> viableBackends,
         ExchangeInfo exchangeInfo,
         RelDataType overrideRowType
+    ) {
+        this(cluster, traitSet, input, viableBackends, exchangeInfo, overrideRowType, null);
+    }
+
+    /**
+     * Full constructor. {@code overrideStorage}, when non-null, must align 1:1 with
+     * {@code overrideRowType} and is returned verbatim by {@link #getOutputFieldStorage()}. QTF
+     * uses it to pair the {@code ___ugsi} column it declares with a matching {@link FieldStorageInfo}.
+     */
+    public OpenSearchExchangeReducer(
+        RelOptCluster cluster,
+        RelTraitSet traitSet,
+        RelNode input,
+        List<String> viableBackends,
+        ExchangeInfo exchangeInfo,
+        RelDataType overrideRowType,
+        List<FieldStorageInfo> overrideStorage
     ) {
         // ConverterImpl makes this a Calcite-recognized trait converter — inserted by
         // Volcano via OpenSearchDistributionTraitDef.convert when a downstream operator
@@ -79,6 +103,7 @@ public OpenSearchExchangeReducer(
         this.viableBackends = viableBackends;
         this.exchangeInfo = exchangeInfo;
         this.overrideRowType = overrideRowType;
+        this.overrideStorage = overrideStorage;
     }
 
     @Override
@@ -98,6 +123,11 @@ public ExchangeInfo getExchangeInfo() {
 
     @Override
     public List<FieldStorageInfo> getOutputFieldStorage() {
+        // When QTF rebuilds this ER with an overrideRowType (declaring ___ugsi), it supplies the
+        // matching overrideStorage so rowType and FSI stay aligned 1:1. Otherwise delegate to input.
+        if (overrideStorage != null) {
+            return overrideStorage;
+        }
         RelNode input = RelNodeUtils.unwrapHep(getInput());
         if (input instanceof OpenSearchRelNode openSearchInput) {
             return openSearchInput.getOutputFieldStorage();
@@ -107,7 +137,15 @@ public List<FieldStorageInfo> getOutputFieldStorage() {
 
     @Override
     public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
-        return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends, exchangeInfo, overrideRowType);
+        return new OpenSearchExchangeReducer(
+            getCluster(),
+            traitSet,
+            sole(inputs),
+            viableBackends,
+            exchangeInfo,
+            overrideRowType,
+            overrideStorage
+        );
     }
 
     /**
@@ -122,7 +160,8 @@ public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
     @Override
     public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
         double rows = mq.getRowCount(getInput());
-        return planner.getCostFactory().makeCost(SETUP_COST + rows, SETUP_COST + rows, 0);
+        double widthFactor = getRowType().getFieldCount();
+        return planner.getCostFactory().makeCost(SETUP_COST + rows * widthFactor, SETUP_COST + rows * widthFactor, 0);
     }
 
     @Override
@@ -138,7 +177,8 @@ public RelNode copyResolved(String backend, List<RelNode> children, List<Operato
             children.getFirst(),
             List.of(backend),
             exchangeInfo,
-            overrideRowType
+            overrideRowType,
+            overrideStorage
         );
     }
 
@@ -151,7 +191,8 @@ public RelNode stripAnnotations(List<RelNode> strippedChildren) {
             strippedChildren.getFirst(),
             viableBackends,
             exchangeInfo,
-            overrideRowType
+            overrideRowType,
+            overrideStorage
         );
     }
 }
@@ -67,6 +67,16 @@ public class OpenSearchLateMaterialization extends SingleRel implements OpenSear
     /** Helper columns consumed internally by the Scatter-Gather stage (not in wrapper output). */
     public static final Set<String> RESERVED_LATE_MATERIALIZATION_FIELDS = Set.of(ROW_ID_FIELD, UGSI_FIELD);
 
+    /**
+     * Trailing-helper layout contract shared by the QTF rewriter (declares these columns) and the
+     * execution side (Stitcher / LM-stage drain read them by name). Helpers always TRAIL the
+     * reduce-set in this order: {@code [..reduce-set.., ___row_id, ___ugsi]}. {@code ___row_id} is
+     * added on the narrowed Scan; {@code ___ugsi} is appended by the ExchangeReducer. Any operator
+     * between the Scan and the wrapper must preserve every trailing helper present in its input, in
+     * this order — never assume a fixed position.
+     */
+    public static final List<String> TRAILING_HELPERS_IN_ORDER = List.of(ROW_ID_FIELD, UGSI_FIELD);
+
     private final List<RelDataTypeField> aboveAnchorPhysicalFields;
     private final List<FieldStorageInfo> aboveAnchorPhysicalFieldStorage;
     private final List<String> viableBackends;

@@ -45,23 +45,51 @@ public class OpenSearchProject extends Project implements OpenSearchRelNode {
 
     private final List<String> viableBackends;
 
+    /**
+     * When true, this Project must stay ABOVE the ExchangeReducer (in the coordinator fragment) —
+     * {@link #computeSelfCost} returns infinite cost unless its input is already gathered
+     * (SINGLETON/ANY), forcing Volcano to place an ER below it. Used to keep an aggregate's literal
+     * config arg (e.g. percentile's {@code 50}) adjacent to the aggregate while a duplicate,
+     * unpinned, physical-only Project pushes below the gather for projection-pushdown. Mirrors the
+     * RexOver gate, which has the same coordinator-side requirement.
+     */
+    private final boolean pinAboveExchange;
+
     public OpenSearchProject(
         RelOptCluster cluster,
         RelTraitSet traitSet,
         RelNode input,
         List<? extends RexNode> projects,
         RelDataType rowType,
         List<String> viableBackends
+    ) {
+        this(cluster, traitSet, input, projects, rowType, viableBackends, false);
+    }
+
+    public OpenSearchProject(
+        RelOptCluster cluster,
+        RelTraitSet traitSet,
+        RelNode input,
+        List<? extends RexNode> projects,
+        RelDataType rowType,
+        List<String> viableBackends,
+        boolean pinAboveExchange
     ) {
         super(cluster, traitSet, List.of(), input, projects, rowType);
         this.viableBackends = viableBackends;
+        this.pinAboveExchange = pinAboveExchange;
     }
 
     @Override
     public List<String> getViableBackends() {
         return viableBackends;
     }
 
+    /** See {@link #pinAboveExchange}. */
+    public boolean isPinAboveExchange() {
+        return pinAboveExchange;
+    }
+
     @Override
     public List<FieldStorageInfo> getOutputFieldStorage() {
         RelNode input = RelNodeUtils.unwrapHep(getInput());
@@ -86,19 +114,21 @@ public List<FieldStorageInfo> getOutputFieldStorage() {
 
     @Override
     public Project copy(RelTraitSet traitSet, RelNode input, List<RexNode> projects, RelDataType rowType) {
-        return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends);
+        return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends, pinAboveExchange);
     }
 
     /**
      * Projects containing {@code RexOver} (window functions) need fully-gathered input so the
-     * window's global frame semantics are correct — infinite cost unless input is SINGLETON.
-     * Volcano picks the plan where an ER sits under this project.
+     * window's global frame semantics are correct. Projects flagged {@link #pinAboveExchange} must
+     * likewise stay in the coordinator fragment (they carry an aggregate's literal config arg). Both
+     * return infinite cost unless input is SINGLETON/ANY — Volcano then picks the plan where an ER
+     * sits under this project.
      *
-     * <p>Plain projects (no RexOver) have no ordering requirement — tiny cost unconditionally.
+     * <p>Plain projects (neither) have no ordering requirement — tiny cost unconditionally.
      */
     @Override
     public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
-        if (!containsOver()) {
+        if (!containsOver() && !pinAboveExchange) {
             return planner.getCostFactory().makeTinyCost();
         }
         // containsOver() is Calcite's own — inherited from Project.
@@ -143,7 +173,15 @@ public RelNode copyResolved(String backend, List<RelNode> children, List<Operato
                 resolvedExprs.add(expr);
             }
         }
-        return new OpenSearchProject(getCluster(), getTraitSet(), children.getFirst(), resolvedExprs, getRowType(), List.of(backend));
+        return new OpenSearchProject(
+            getCluster(),
+            getTraitSet(),
+            children.getFirst(),
+            resolvedExprs,
+            getRowType(),
+            List.of(backend),
+            pinAboveExchange
+        );
     }
 
     @Override