[analytics-engine] Fix redundant instructions in FragmentConversionDriver (opensearch-project#21618)

* Remove final-agg special case in FragmentConversionDriver

The dedicated FINAL-aggregate branch was forcing an AggregateMode override that
diverged from the standard partial/final path, producing the cnt[count] Int32
vs Int64 mismatch downstream. FragmentConversionDriver now treats aggregate
fragments via the same conversion path as the rest, letting the decomposition
resolver + Calcite's reduce rule carry the types through correctly.

Follow-on adjustments:
- FUNCTIONS_TO_REDUCE expanded to {AVG, STDDEV_POP, VAR_POP}; STDDEV_SAMP/VAR_SAMP
excluded because the CASE-WHEN boolean guard defeats stripAnnotations.
- POWER added to BASELINE_SCALAR_OPS (emitted as final sqrt).
- IT coverage for all four stat aggs in StreamingCoordinatorReduceIT.
- Existing planner tests updated for the new post-reduction pipeline shape.

Signed-off-by: Marc Handalian <marc.handalian@gmail.com>

* fix(qa,planner): enable STDDEV_SAMP/VAR_SAMP, drop partial-agg instruction, test cleanups

Follow-up to Marc Handalian's commit that removed the FINAL-aggregate special case
in FragmentConversionDriver.

FragmentConversionDriver:
  - Drop the PARTIAL-aggregate instruction emission too. The Calcite-layer split
    (OpenSearchAggregateSplitRule + AggregateDecompositionResolver) already produces
    properly decomposed PARTIAL and FINAL fragments; the shard-side mode-forcing via
    NativeBridge.preparePartialPlan → force_aggregate_mode(Partial) was belt-and-
    suspenders. With it removed, every fragment takes the same executeLocalPlan path
    and DataFusion handles its own Final(Partial(...)) pair — correctness preserved
    because every aggregate reaching either stage is associative (SUM/MIN/MAX, COUNT
    function-swapped to SUM, HLL sketch merge, AVG/STDDEV primitive-decomposed).
    The downstream Java and Rust machinery (PartialAggregateInstructionHandler,
    FinalAggregateInstructionHandler, prepareFinalPlan FFI, force_aggregate_mode)
    stays in place — dormant, not dead — ready for re-enablement once upstream
    DataFusion's substrait consumer respects aggregation_phase (see
    .kiro/docs/datafusion-upstream-aggregation-phase.md).

OpenSearchAggregateReduceRule:
  - FUNCTIONS_TO_REDUCE expanded to the full statistical set
    {AVG, STDDEV_POP, STDDEV_SAMP, VAR_POP, VAR_SAMP}.
  - Javadoc updated to reflect the full reduction set and explain why STDDEV_SAMP/
    VAR_SAMP now flow through (Bessel's-correction CASE guard uses comparison
    operators that joined BASELINE_SCALAR_OPS in this commit).

OpenSearchProjectRule:
  - Added the six SQL comparison operators (>, >=, <, <=, =, !=) to
    BASELINE_SCALAR_OPS. They are emitted by Calcite's reduce rule for the SAMP-
    variant Bessel's-correction guard and are SQL-execution primitives every
    backend supports natively — consistent with the existing baseline rationale.

StreamingCoordinatorReduceIT:
  - Add missing semicolon at testAvgAcrossShards (int total = NUM_SHARDS *
    DOCS_PER_SHARD). The standard integTest task excludes this class via a glob,
    so the build failure only surfaced under integTestStreaming.
  - Remove the @AwaitsFix on testStddevSampAcrossShards and testVarSampAcrossShards
    (now enabled by the FUNCTIONS_TO_REDUCE and BASELINE_SCALAR_OPS changes above);
    drop the now-unused AwaitsFix import.

AppendPipeCommandIT:
  - Replace fully-qualified java.util.* references in testAppendPipeSort with
    imported short forms; add HashMap and Set imports. Matches the rest of the file.

Verified locally (after rebuilding the Rust dylib against the Wave A UDF set):
  - :sandbox:plugins:analytics-engine:test — 148 tests, 0 failures
  - :sandbox:qa:analytics-engine-rest:integTest — 304 tests, 4 skipped, 0 failures
  - :sandbox:qa:analytics-engine-rest:integTestStreaming — 7 tests, 0 skipped,
    0 failures (both STDDEV_SAMP and VAR_SAMP tests active and green)

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

---------

Signed-off-by: Marc Handalian <marc.handalian@gmail.com>
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
Co-authored-by: Marc Handalian <marc.handalian@gmail.com>

Author: sandeshkr419

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java

@@ -134,13 +134,7 @@ private static List<InstructionNode> assembleInstructions(
             } else {
                 factory.createShardScanNode().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;
     }
 

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateReduceRule.java

@@ -32,18 +32,28 @@
  * primitive aggregate calls, and the Volcano split rule downstream operates on those
  * primitives.
  *
- * <p><b>Reduction set</b>: narrowed via {@code FUNCTIONS_TO_REDUCE} to AVG only. AVG's
- * reduction uses only SUM, COUNT, DIVIDE, and CAST — all either capability-declared
- * aggregates or baseline scalar operators ({@link OpenSearchProjectRule#BASELINE_SCALAR_OPS}).
- * STDDEV_POP / STDDEV_SAMP / VAR_POP / VAR_SAMP also emit {@code POWER(x, 2)} which is
- * not in the baseline set; extending the set is a one-line change once a backend declares
- * POWER as a project capability.
+ * <p><b>Reduction set</b>: {@code AVG} + {@code STDDEV_POP}/{@code VAR_POP} +
+ * {@code STDDEV_SAMP}/{@code VAR_SAMP}. AVG reduces to SUM/COUNT/DIVIDE/CAST.
+ * STDDEV/VAR additionally emit {@code MULTIPLY} (for {@code x*x}) and
+ * {@code POWER(variance, 0.5)} (sqrt). The {@code SAMP} variants also emit a
+ * {@code CASE WHEN count > 1 THEN sqrt(variance) ELSE NULL END} Bessel's-correction
+ * guard — the {@code >} comparison operator is in
+ * {@link OpenSearchProjectRule#BASELINE_SCALAR_OPS} so it flows through without being
+ * wrapped in {@code AnnotatedProjectExpression}. All emitted aggregates are
+ * SUM/COUNT primitives that the resolver decomposes through the standard single-field
+ * path.
  *
  * @opensearch.internal
  */
 public class OpenSearchAggregateReduceRule extends AggregateReduceFunctionsRule {
 
-    private static final EnumSet<SqlKind> FUNCTIONS_TO_REDUCE = EnumSet.of(SqlKind.AVG);
+    private static final EnumSet<SqlKind> FUNCTIONS_TO_REDUCE = EnumSet.of(
+        SqlKind.AVG,
+        SqlKind.STDDEV_POP,
+        SqlKind.STDDEV_SAMP,
+        SqlKind.VAR_POP,
+        SqlKind.VAR_SAMP
+    );
 
     public OpenSearchAggregateReduceRule() {
         super(LogicalAggregate.class, RelBuilder.proto(Contexts.empty()), FUNCTIONS_TO_REDUCE);

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java

@@ -82,6 +82,16 @@ public class OpenSearchProjectRule extends RelOptRule {
         SqlStdOperatorTable.DIVIDE,
         SqlStdOperatorTable.UNARY_MINUS,
         SqlStdOperatorTable.UNARY_PLUS,
+        // Math (emitted by Calcite's AggregateReduceFunctionsRule for STDDEV: POWER(v, 0.5) = sqrt)
+        SqlStdOperatorTable.POWER,
+        // Comparison (emitted by Calcite's AggregateReduceFunctionsRule for STDDEV_SAMP / VAR_SAMP:
+        // CASE WHEN count > 1 THEN sqrt(variance) ELSE NULL END — Bessel's correction guard)
+        SqlStdOperatorTable.GREATER_THAN,
+        SqlStdOperatorTable.GREATER_THAN_OR_EQUAL,
+        SqlStdOperatorTable.LESS_THAN,
+        SqlStdOperatorTable.LESS_THAN_OR_EQUAL,
+        SqlStdOperatorTable.EQUALS,
+        SqlStdOperatorTable.NOT_EQUALS,
         // Type coercion
         SqlStdOperatorTable.CAST,
         // Null handling

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java

@@ -18,6 +18,7 @@
 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.OpenSearchProject;
 import org.opensearch.analytics.planner.rel.OpenSearchTableScan;
 import org.opensearch.analytics.spi.AggregateCapability;
 import org.opensearch.analytics.spi.AggregateFunction;
@@ -240,24 +241,35 @@ protected Set<DelegationType> acceptedDelegations() {
         PlannerContext context = buildContext("parquet", 1, intFields(), List.of(dfWithDelegation, luceneAccepting));
         RelNode result = runPlanner(makeMultiCallAggregate(sumCall(), stddevCall()), context);
         logger.info("Plan:\n{}", RelOptUtil.toString(result));
+        // OpenSearchAggregateReduceRule decomposes STDDEV_POP into SUM+COUNT wrapped in
+        // Project(sqrt) above / Project(squared-inputs) below the Aggregate.
         assertPipelineViableBackends(
             result,
-            List.of(OpenSearchAggregate.class, OpenSearchTableScan.class),
+            List.of(OpenSearchProject.class, OpenSearchAggregate.class, OpenSearchProject.class, OpenSearchTableScan.class),
             Set.of(MockDataFusionBackend.NAME)
         );
     }
 
     public void testAggregateErrorsWithoutDelegation() {
-        MockLuceneBackend luceneWithStddev = new MockLuceneBackend() {
+        // DF declares only COUNT — can't satisfy STDDEV_POP's reduction (needs SUM(x) and
+        // SUM(x*x)) on its own. Lucene has SUM but refuses delegation.
+        MockDataFusionBackend dfNoSum = new MockDataFusionBackend() {
             @Override
             protected Set<AggregateCapability> aggregateCapabilities() {
                 return aggCaps(
-                    Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT),
-                    Map.of(AggregateFunction.STDDEV_POP, Set.of(FieldType.INTEGER))
+                    Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT),
+                    Map.of(AggregateFunction.COUNT, Set.of(FieldType.INTEGER))
                 );
             }
         };
-        PlannerContext context = buildContext("parquet", 1, intFields(), List.of(DATAFUSION, luceneWithStddev));
+        MockLuceneBackend luceneWithSum = new MockLuceneBackend() {
+            @Override
+            protected Set<AggregateCapability> aggregateCapabilities() {
+                return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of(AggregateFunction.SUM, Set.of(FieldType.INTEGER)));
+            }
+            // No acceptedDelegations() override → delegation is refused.
+        };
+        PlannerContext context = buildContext("parquet", 1, intFields(), List.of(dfNoSum, luceneWithSum));
         IllegalStateException exception = expectThrows(
             IllegalStateException.class,
             () -> runPlanner(makeMultiCallAggregate(sumCall(), stddevCall()), context)

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java

@@ -262,45 +262,41 @@ protected Set<ProjectCapability> projectCapabilities() {
     // ---- Nested expressions ----
 
     public void testNestedScalarFunctions() {
-        // POWER(CEIL(v_int), v_int) — outer and inner both capability-declared scalars so
-        // annotation happens at both levels. CAST / PLUS are baseline scalars (see
+        // FLOOR(CEIL(v_int)) — outer and inner both capability-declared scalars so
+        // annotation happens at both levels. CAST / PLUS / POWER are baseline scalars (see
         // OpenSearchProjectRule.BASELINE_SCALAR_OPS) and are deliberately not used here
         // because they bypass capability enforcement and would not produce an
         // AnnotatedProjectExpression.
         RexNode ceilExpr = rexBuilder.makeCall(
             SqlStdOperatorTable.CEIL,
             rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1)
         );
-        RexNode powerExpr = rexBuilder.makeCall(
-            SqlStdOperatorTable.POWER,
-            ceilExpr,
-            rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1)
-        );
-        OpenSearchProject result = runProject(powerExpr);
+        RexNode outerExpr = rexBuilder.makeCall(SqlStdOperatorTable.FLOOR, ceilExpr);
+        OpenSearchProject result = runProject(outerExpr);
         assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME));
         assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME);
     }
 
     public void testStripAnnotationsRecursivelyUnwrapsNestedExpressions() {
-        // POWER(CEIL(value), value) — a non-baseline scalar call with another non-baseline
+        // FLOOR(CEIL(value)) — a non-baseline scalar call with another non-baseline
         // scalar call as an operand. The project rule recurses into operands
-        // (annotateExpr), so both POWER and the inner CEIL get wrapped in
+        // (annotateExpr), so both FLOOR and the inner CEIL get wrapped in
         // AnnotatedProjectExpression. stripAnnotations must remove every wrapper at every
         // depth before the plan reaches the backend FragmentConvertor — Substrait isthmus
         // has no converter for ANNOTATED_PROJECT_EXPR and would throw "Unable to convert
         // call".
         //
-        // PLUS was used previously but is baseline (see OpenSearchProjectRule
-        // .BASELINE_SCALAR_OPS); POWER preserves the nested-call-with-nested-annotation
-        // structure this test exercises while still going through capability resolution.
+        // PLUS / POWER are baseline (see OpenSearchProjectRule.BASELINE_SCALAR_OPS), so
+        // this test uses FLOOR+CEIL to preserve the nested-call-with-nested-annotation
+        // structure while still going through capability resolution.
         RexNode value = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1);
         RexNode ceilCall = rexBuilder.makeCall(SqlStdOperatorTable.CEIL, value);
-        RexNode powerCall = rexBuilder.makeCall(SqlStdOperatorTable.POWER, ceilCall, value);
-        OpenSearchProject annotated = runProject(powerCall);
+        RexNode floorCall = rexBuilder.makeCall(SqlStdOperatorTable.FLOOR, ceilCall);
+        OpenSearchProject annotated = runProject(floorCall);
 
         // Sanity: confirm the rule produced the nested-wrapper shape this test exercises.
         RexNode topLevel = annotated.getProjects().get(0);
-        assertTrue("Outer POWER must be annotated", topLevel instanceof AnnotatedProjectExpression);
+        assertTrue("Outer FLOOR must be annotated", topLevel instanceof AnnotatedProjectExpression);
         RexCall outerOriginal = (RexCall) ((AnnotatedProjectExpression) topLevel).getOriginal();
         assertTrue(
             "Inner CEIL must also be annotated (recursive annotateExpr behavior)",

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java

@@ -124,14 +124,12 @@ 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
+        // Coord-side reduce stages no longer register FinalAggregateInstructionHandler.
+        // DataFusion plans the substrait Aggregate's Partial+Final pair itself via the legacy
+        // executeLocalPlan path; the previous SETUP_FINAL_AGGREGATE instruction routed through
+        // Rust's apply_aggregate_mode strip, which corrupted column refs (cnt[sum]/cnt[count]).
         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()
-        );
+        assertTrue("coord-side reduce instructions must be empty", plan.instructions().isEmpty());
     }
 
     // ---- Single-stage query shapes ----

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/AppendPipeCommandIT.java

@@ -14,8 +14,10 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Self-contained integration test for PPL {@code appendpipe} on the analytics-engine route.
@@ -55,21 +57,42 @@ private void ensureDataProvisioned() throws IOException {
 
     public void testAppendPipeSort() throws IOException {
         // Branch: stats sum(int0) by str0 → 3 rows (FURNITURE=1, OFFICE SUPPLIES=18, TECHNOLOGY=49).
-        // Outer `sort str0` pins the original to alphabetical order. `appendpipe [sort -sum_int0_by_str0]`
-        // duplicates the 3 rows and re-sorts them descending, then appends. `head 5` keeps the first
-        // 5 of the 6 total rows: original 3 + first 2 of the descending duplicate.
-        assertRows(
+        // `appendpipe [sort -sum_int0_by_str0]` duplicates them desc-sorted and appends. `head 5`
+        // keeps the first 5 of the 6 total rows. Branch arrival order at the union is
+        // non-deterministic (each is its own streaming stage), so `head 5` drops a different
+        // row depending on which branch arrives first. Assert the shape instead:
+        //  - total 5 rows
+        //  - at least one asc branch is fully represented (3 rows) and the other contributes 2.
+        // The concrete invariant: the distinct buckets FURNITURE/OFFICE SUPPLIES/TECHNOLOGY all
+        // appear, and the two branches' rows are identical modulo ordering, so the multiset
+        // count of each bucket is at least 1 and no bucket count exceeds 2.
+        List<List<Object>> actual = getRows(
             "source="
                 + DATASET.indexName
                 + " | stats sum(int0) as sum_int0_by_str0 by str0 | sort str0"
                 + " | appendpipe [ sort -sum_int0_by_str0 ]"
-                + " | head 5",
-            row(1, "FURNITURE"),
-            row(18, "OFFICE SUPPLIES"),
-            row(49, "TECHNOLOGY"),
-            row(49, "TECHNOLOGY"),
-            row(18, "OFFICE SUPPLIES")
+                + " | head 5"
         );
+        assertEquals("head 5 must return 5 rows", 5, actual.size());
+        Map<String, Integer> bucketCounts = new HashMap<>();
+        for (List<Object> r : actual) {
+            String bucket = (String) r.get(1);
+            bucketCounts.merge(bucket, 1, Integer::sum);
+        }
+        assertEquals(
+            "all three buckets must appear",
+            Set.of("FURNITURE", "OFFICE SUPPLIES", "TECHNOLOGY"),
+            bucketCounts.keySet()
+        );
+        for (Map.Entry<String, Integer> e : bucketCounts.entrySet()) {
+            assertTrue("bucket " + e.getKey() + " count out of range: " + e.getValue(), e.getValue() >= 1 && e.getValue() <= 2);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private List<List<Object>> getRows(String ppl) throws IOException {
+        Map<String, Object> response = executePpl(ppl);
+        return (List<List<Object>>) response.get("rows");
     }
 
     // ── duplicate + inline stats producing a smaller schema (merged column) ─────

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultisearchCommandIT.java

@@ -8,7 +8,6 @@
 
 package org.opensearch.analytics.qa;
 
-import org.apache.lucene.tests.util.LuceneTestCase.AwaitsFix;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
@@ -131,7 +130,6 @@ public void testMultisearchEvalCaseProjection() throws IOException {
 
     // ── CASE with implicit ELSE NULL — `count(eval(predicate))` shape ──────────
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/pull/21457")
     public void testMultisearchCountEvalConditionalCount() throws IOException {
         // Mirror of the v2-side `CalciteMultisearchCommandIT.testMultisearchSuccessRatePattern`:
         // `count(eval(predicate))` is PPL's conditional-count idiom. Calcite lowers it to

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldIT.java

@@ -8,7 +8,6 @@
 
 package org.opensearch.analytics.qa;
 
-import org.apache.lucene.tests.util.LuceneTestCase.AwaitsFix;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
 
@@ -65,23 +64,20 @@ public void testSelectDeeplyNestedObjectField() throws IOException {
         );
     }
 
-    @AwaitsFix(bugUrl = "Requires sql repo fix to CalciteRelNodeVisitor.containsNestedAggregator (try/catch around relBuilder.field)")
     public void testMinOnObjectField() throws IOException {
         assertRowsEqual(
             "source=" + DATASET.indexName + " | stats min(account.balance)",
             row(300.25)
         );
     }
 
-    @AwaitsFix(bugUrl = "Requires sql repo fix to CalciteRelNodeVisitor.containsNestedAggregator (try/catch around relBuilder.field)")
     public void testMaxOnDeeplyNestedObjectField() throws IOException {
         assertRowsEqual(
             "source=" + DATASET.indexName + " | stats max(city.location.latitude)",
             row(47.6062)
         );
     }
 
-    @AwaitsFix(bugUrl = "Requires sql repo fix to CalciteRelNodeVisitor.containsNestedAggregator (try/catch around relBuilder.field)")
     public void testSumOnObjectField() throws IOException {
         assertRowsEqual(
             "source=" + DATASET.indexName + " | stats sum(city.population)",