[Analytics-Engine] Fix multi-shard aggregate split for non-prefix group keys (#22065)

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

Author: mch2

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java

@@ -120,7 +120,15 @@ public OpenSearchAggregate(
         Map<Integer, List<RexLiteral>> finalExtraLiteralArgs,
         List<IntermediateField> perCallIntermediateField
     ) {
-        super(cluster, traitSet, List.of(), input, groupSet, groupSets, aggCalls);
+        super(
+            cluster,
+            traitSet,
+            List.of(),
+            input,
+            frontGroupSetForFinal(groupSet, mode),
+            frontGroupSetsForFinal(groupSet, groupSets, mode),
+            aggCalls
+        );
         this.mode = mode;
         this.viableBackends = viableBackends;
         this.callAnnotations = Map.copyOf(callAnnotations);
@@ -129,6 +137,30 @@ public OpenSearchAggregate(
         this.perCallIntermediateField = Collections.unmodifiableList(new ArrayList<>(perCallIntermediateField));
     }
 
+    /**
+     * FINAL reads PARTIAL's output, where Calcite has fronted the group keys to {@code 0..n-1}; group
+     * on the prefix range so a non-prefix key (e.g. {@code avg(x) by span(y,5)} → {@code {1}}) isn't
+     * read as an agg-state column. No-op for SINGLE/PARTIAL (raw input) and already-fronted sets.
+     */
+    private static ImmutableBitSet frontGroupSetForFinal(ImmutableBitSet groupSet, AggregateMode mode) {
+        return mode == AggregateMode.FINAL ? ImmutableBitSet.range(groupSet.cardinality()) : groupSet;
+    }
+
+    /**
+     * Keeps {@code groupSets} consistent with the fronted {@code groupSet}. PPL only emits simple
+     * (single-set) aggregates, so this collapses to one set; revisit if GROUPING SETS is ever added.
+     */
+    private static List<ImmutableBitSet> frontGroupSetsForFinal(
+        ImmutableBitSet groupSet,
+        List<ImmutableBitSet> groupSets,
+        AggregateMode mode
+    ) {
+        if (mode != AggregateMode.FINAL || groupSets == null || groupSets.isEmpty()) {
+            return groupSets;
+        }
+        return List.of(ImmutableBitSet.range(groupSet.cardinality()));
+    }
+
     /** Builds a FINAL aggregate post-rewrite; clears both stashes so a later {@code copy()} can't replay them. */
     public static OpenSearchAggregate finalAfterRewrite(OpenSearchAggregate prior, RelNode newInput, List<AggregateCall> rebuiltCalls) {
         return new OpenSearchAggregate(

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

@@ -9,8 +9,13 @@
 package org.opensearch.analytics.planner;
 
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.ImmutableBitSet;
 
+import java.util.List;
+
 /**
  * Plan-shape tests for {@link org.opensearch.analytics.planner.rel.OpenSearchAggregate}.
  *
@@ -20,6 +25,49 @@
  */
 public class AggregatePlanShapeTests extends PlanShapeTestBase {
 
+    // ---- Non-prefix groupSet: agg arg projected BEFORE the group key (the `avg(x) by span(y,5)` ----
+    // ---- shape). PARTIAL fronts the key per Calcite's contract, so FINAL must group on the prefix ----
+    // ---- range {0}, NOT the original {1}. Regression guard for the multi-shard span bug where FINAL ----
+    // ---- grouped on the partial-state column and produced spurious un-coalesced buckets. ----
+
+    /** Builds sum(col0) grouped by col1 → non-prefix groupSet {1}. */
+    private RelNode nonPrefixGroupedSum() {
+        RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+        AggregateCall sum = AggregateCall.create(
+            SqlStdOperatorTable.SUM,
+            false,
+            List.of(0),
+            -1,
+            scan,
+            typeFactory.createSqlType(SqlTypeName.INTEGER),
+            "total_status"
+        );
+        return makeAggregate(scan, ImmutableBitSet.of(1), sum);
+    }
+
+    public void testNonPrefixGroupSet_1shard() {
+        RelNode result = runPlanner(nonPrefixGroupedSum(), singleShardContext());
+        assertPlanShape("""
+            OpenSearchAggregate(group=[{1}], total_status=[SUM($0)], mode=[SINGLE], viableBackends=[[mock-parquet]])
+              OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]])
+            """, result);
+    }
+
+    public void testNonPrefixGroupSet_2shard() {
+        // FINAL groups on {0} (key fronted by PARTIAL), NOT the original {1}. Before the fix FINAL
+        // carried {1} and grouped on the SUM column instead of the key.
+        RelNode result = runPlanner(nonPrefixGroupedSum(), multiShardContext());
+        assertPlanShape(
+            """
+                OpenSearchAggregate(group=[{0}], total_status=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]])
+                  OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]])
+                    OpenSearchAggregate(group=[{1}], total_status=[SUM($0)], mode=[PARTIAL], viableBackends=[[mock-parquet]])
+                      OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]])
+                """,
+            result
+        );
+    }
+
     public void testStatsCountStarByKey_1shard() {
         RelNode scan = stubScan(mockTable("test_index", "status", "size"));
         RelNode plan = makeAggregate(scan, countStarCall(scan));

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

@@ -327,7 +327,7 @@ public void testJoinWithDifferentGroupKeys_multiShard() {
                     OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]])
                       OpenSearchAggregate(group=[{0}], s=[SUM($1)], mode=[PARTIAL], viableBackends=[[mock-parquet]])
                         OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]])
-                  OpenSearchAggregate(group=[{1}], s=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]])
+                  OpenSearchAggregate(group=[{0}], s=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]])
                     OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]])
                       OpenSearchAggregate(group=[{1}], s=[SUM($0)], mode=[PARTIAL], viableBackends=[[mock-parquet]])
                         OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]])

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

@@ -16,6 +16,16 @@
  * stddev/var, span, values/list) and approximate tier {@code approx/} (distinct_count/dc/percentile/
  * median). {@code distinct_count}/{@code dc} are rewritten to {@code APPROX_COUNT_DISTINCT} and
  * computed once at the coordinator, so per-shard distinct counts are never summed.
+ *
+ * <p>The {@code agg/span_*} set covers every aggregate over a {@code span(...)} group key — a
+ * non-prefix groupSet that exercises the PARTIAL/FINAL key-fronting fix. A few of those
+ * ({@code span_distinct_count_label}, {@code span_percentile_50}, {@code span_median}) use
+ * approximate algorithms (HLL / percentile_approx) yet sit in the EXACT tier: on the fixed 30-row
+ * {@code merge_coverage} dataset they are below the approximation threshold and return exact,
+ * shard-invariant results, so the 1-shard==2-shard differential holds. If that dataset ever grows
+ * past the approximation threshold these must move to the {@code approx/} tier — but note the
+ * approx tier's tolerance aligner keys rows on non-numeric group cells, so a numeric {@code span}
+ * key can't be aligned there without a harness change.
  */
 public class TwoShardAggregationIT extends TwoShardReduceTestCase {
 

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

@@ -31,7 +31,6 @@ protected Map<String, String> knownIssues() {
         Map<String, String> m = new LinkedHashMap<>();
         m.put("cmd_search", "search numeric comparison lowers to a Lucene query_string that matches zero docs on numeric fields (wrong result)");
         m.put("cmd_appendcols", "not in the PPL grammar (SyntaxCheckException)");
-        m.put("cmd_timechart", "requires an @timestamp default field");
         return m;
     }
 }

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

@@ -0,0 +1,45 @@
+/*
+ * 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.qa;
+
+import org.opensearch.client.Request;
+
+import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * Same exact {@code agg/} differential + golden suite as {@link TwoShardAggregationIT}, but with
+ * {@code analytics.shard_bucket_oversampling_factor} turned on. Oversampling changes how many
+ * candidate buckets each shard ships to the coordinator (and inserts the TopK shard-limit rewrite),
+ * so this exercises that rewrite composing with the non-prefix-groupSet span aggregates (and with
+ * the sort/top-N-over-span shapes). At this dataset scale oversampling does not change the final
+ * exact result, so the 1-shard == 2-shard differential and the goldens must still hold — proving
+ * the oversampling/TopK rewrite doesn't perturb exact aggregation over a fronted-key FINAL stage.
+ *
+ * <p>Only the exact {@code agg} tier is run — the {@code approx} tier is genuinely non-deterministic
+ * under oversampling and is covered without oversampling by {@link TwoShardAggregationIT}.
+ */
+public class TwoShardOversamplingAggregationIT extends TwoShardReduceTestCase {
+
+    @Override
+    protected Map<String, Boolean> tiers() {
+        Map<String, Boolean> t = new LinkedHashMap<>();
+        t.put("agg", false); // exact tier only — oversampling must be a no-op here
+        return t;
+    }
+
+    @Override
+    protected void onBeforeQuery() throws IOException {
+        super.onBeforeQuery();
+        Request req = new Request("PUT", "/_cluster/settings");
+        req.setJsonEntity("{\"persistent\":{\"analytics.shard_bucket_oversampling_factor\": 2.0}}");
+        client().performRequest(req);
+    }
+}

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

@@ -158,9 +158,17 @@ private String runDifferential(String dir, String name) {
             Map<String, Object> golden = loadGolden(queryDir + "/expected/" + name + ".json");
             if (golden != null) {
                 normalizeArrayCells(golden);
-                String pin = ResponseValidator.compareData(golden, r2, name + " [2-shard vs golden]");
-                if (pin != null) {
-                    return pin;
+                // Pin BOTH shard counts to the golden. The differential above already proves
+                // r1 == r2, so 2-shard-vs-golden alone would catch a regression transitively — but
+                // pinning 1-shard directly guards the baseline even if the differential is ever
+                // relaxed, and a baseline-vs-coordinator regression surfaces on the right side.
+                String pin1 = ResponseValidator.compareData(golden, r1, name + " [1-shard vs golden]");
+                if (pin1 != null) {
+                    return pin1;
+                }
+                String pin2 = ResponseValidator.compareData(golden, r2, name + " [2-shard vs golden]");
+                if (pin2 != null) {
+                    return pin2;
                 }
             }
             return null;

sandbox/qa/analytics-engine-rest/src/test/resources/datasets/merge_coverage/agg/expected/span_avg_double.json

@@ -0,0 +1,32 @@
+{
+  "rows": [
+    [
+      3.75,
+      0
+    ],
+    [
+      10.5,
+      5
+    ],
+    [
+      18.0,
+      10
+    ],
+    [
+      25.5,
+      15
+    ],
+    [
+      33.0,
+      20
+    ],
+    [
+      40.5,
+      25
+    ],
+    [
+      45.0,
+      30
+    ]
+  ]
+}

sandbox/qa/analytics-engine-rest/src/test/resources/datasets/merge_coverage/agg/expected/span_count.json

@@ -0,0 +1,32 @@
+{
+  "rows": [
+    [
+      4,
+      0
+    ],
+    [
+      5,
+      5
+    ],
+    [
+      5,
+      10
+    ],
+    [
+      5,
+      15
+    ],
+    [
+      5,
+      20
+    ],
+    [
+      5,
+      25
+    ],
+    [
+      1,
+      30
+    ]
+  ]
+}

sandbox/qa/analytics-engine-rest/src/test/resources/datasets/merge_coverage/agg/expected/span_count_avg.json

@@ -0,0 +1,39 @@
+{
+  "rows": [
+    [
+      4,
+      3.75,
+      0
+    ],
+    [
+      5,
+      10.5,
+      5
+    ],
+    [
+      5,
+      18.0,
+      10
+    ],
+    [
+      5,
+      25.5,
+      15
+    ],
+    [
+      5,
+      33.0,
+      20
+    ],
+    [
+      5,
+      40.5,
+      25
+    ],
+    [
+      1,
+      45.0,
+      30
+    ]
+  ]
+}