Add EXTENDED_STATS bundle aggregate decomposition

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java

@@ -47,6 +47,11 @@ public enum AggregateFunction {
     // so resolution goes via name lookup in fromNameOrError.
     STATS(Type.STATE_EXPANDING, SqlKind.OTHER),
 
+    // EXTENDED_STATS — superset of STATS adding sum_of_squares, variance/std_deviation
+    // triplets, and a nested std_deviation_bounds struct. Decomposition is handled by
+    // OpenSearchExtendedStatsReduceRule during HEP, same model as STATS.
+    EXTENDED_STATS(Type.STATE_EXPANDING, SqlKind.OTHER),
+
     // Statistical — fixed-size state, multi-pass or running stats. Handled by
     // OpenSearchAggregateReduceRule (once FUNCTIONS_TO_REDUCE is extended to include them)
     // — no intermediateFields here either.
@@ -211,6 +216,7 @@ public SqlAggFunction toSqlAggFunction() {
             case AVG -> SqlStdOperatorTable.AVG;
             case APPROX_COUNT_DISTINCT -> SqlStdOperatorTable.APPROX_COUNT_DISTINCT;
             case STATS -> OpenSearchAggregateOperators.STATS;
+            case EXTENDED_STATS -> OpenSearchAggregateOperators.EXTENDED_STATS;
             default -> throw new IllegalStateException("No SqlAggFunction mapping for: " + this);
         };
     }

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/OpenSearchAggregateOperators.java

@@ -31,6 +31,12 @@
  *       Decomposed at HEP planning time by
  *       {@code OpenSearchStatsReduceRule} into primitive COUNT/MIN/MAX/SUM aggregate calls
  *       plus a Project that builds the struct, so no executor ever sees STATS.</li>
+ *   <li>{@link #EXTENDED_STATS} — superset of {@link #STATS}, returning a 13-field struct
+ *       that adds {@code sum_of_squares}, the {@code variance}/{@code std_deviation}
+ *       triplets, and the nested 6-field {@code std_deviation_bounds} struct. Decomposed
+ *       at HEP planning time by {@code OpenSearchExtendedStatsReduceRule} into the same
+ *       four primitives as STATS plus {@code SUM(x*x)} (computed via a Project beneath the
+ *       rebuilt aggregate), with the variance/stddev/bounds derived in the assembly Project.</li>
  * </ul>
  *
  * @opensearch.internal
@@ -50,6 +56,34 @@ private OpenSearchAggregateOperators() {}
     public static final String STATS_FIELD_SUM = "sum";
     public static final String STATS_FIELD_AVG = "avg";
 
+    /**
+     * Field names of the struct returned by {@link #EXTENDED_STATS}. Order matches legacy
+     * OpenSearch {@code InternalExtendedStats.toXContent}: count/min/max/avg/sum from
+     * {@code InternalStats}, then sum_of_squares, variance triplet, std_deviation triplet,
+     * and the nested {@code std_deviation_bounds} struct (with the six bound sub-fields).
+     */
+    public static final String EXTENDED_STATS_FIELD_SUM_OF_SQUARES = "sum_of_squares";
+    public static final String EXTENDED_STATS_FIELD_VARIANCE = "variance";
+    public static final String EXTENDED_STATS_FIELD_VARIANCE_POPULATION = "variance_population";
+    public static final String EXTENDED_STATS_FIELD_VARIANCE_SAMPLING = "variance_sampling";
+    public static final String EXTENDED_STATS_FIELD_STD_DEVIATION = "std_deviation";
+    public static final String EXTENDED_STATS_FIELD_STD_DEVIATION_POPULATION = "std_deviation_population";
+    public static final String EXTENDED_STATS_FIELD_STD_DEVIATION_SAMPLING = "std_deviation_sampling";
+    public static final String EXTENDED_STATS_FIELD_STD_DEVIATION_BOUNDS = "std_deviation_bounds";
+    public static final String EXTENDED_STATS_BOUND_UPPER = "upper";
+    public static final String EXTENDED_STATS_BOUND_LOWER = "lower";
+    public static final String EXTENDED_STATS_BOUND_UPPER_POPULATION = "upper_population";
+    public static final String EXTENDED_STATS_BOUND_LOWER_POPULATION = "lower_population";
+    public static final String EXTENDED_STATS_BOUND_UPPER_SAMPLING = "upper_sampling";
+    public static final String EXTENDED_STATS_BOUND_LOWER_SAMPLING = "lower_sampling";
+
+    /**
+     * Sigma multiplier used when computing {@code std_deviation_bounds}. Hardcoded to match
+     * the legacy {@code InternalExtendedStats} default; the bounds are
+     * {@code avg ± SIGMA * std_deviation}.
+     */
+    public static final double EXTENDED_STATS_DEFAULT_SIGMA = 2.0;
+
     /**
      * Return-type inference for {@link #STATS}. Produces a struct whose component types match
      * what Calcite's built-in COUNT/MIN/MAX/SUM aggregates would produce for the same operand,
@@ -58,6 +92,9 @@ private OpenSearchAggregateOperators() {}
      */
     private static final SqlReturnTypeInference STATS_RETURN_TYPE_INFERENCE = OpenSearchAggregateOperators::inferStatsReturnType;
 
+    private static final SqlReturnTypeInference EXTENDED_STATS_RETURN_TYPE_INFERENCE =
+        OpenSearchAggregateOperators::inferExtendedStatsReturnType;
+
     /**
      * STATS bundle aggregate. Decomposed at HEP planning time — never reaches the executor.
      */
@@ -75,6 +112,25 @@ private OpenSearchAggregateOperators() {}
     ) {
     };
 
+    /**
+     * EXTENDED_STATS bundle aggregate. Superset of {@link #STATS} — adds sum_of_squares,
+     * variance/std_deviation triplets, and a nested std_deviation_bounds struct. Decomposed
+     * at HEP planning time — never reaches the executor.
+     */
+    public static final SqlAggFunction EXTENDED_STATS = new SqlAggFunction(
+        "EXTENDED_STATS",
+        null,
+        SqlKind.OTHER,
+        EXTENDED_STATS_RETURN_TYPE_INFERENCE,
+        null,
+        OperandTypes.NUMERIC,
+        SqlFunctionCategory.USER_DEFINED_FUNCTION,
+        false,
+        false,
+        Optionality.FORBIDDEN
+    ) {
+    };
+
     private static RelDataType inferStatsReturnType(SqlOperatorBinding opBinding) {
         RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
         RelDataType inputType = opBinding.getOperandType(0);
@@ -103,4 +159,50 @@ private static RelDataType inferStatsReturnType(SqlOperatorBinding opBinding) {
             .add(STATS_FIELD_SUM, sumType)
             .build();
     }
+
+    private static RelDataType inferExtendedStatsReturnType(SqlOperatorBinding opBinding) {
+        RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
+        RelDataType inputType = opBinding.getOperandType(0);
+
+        RelDataType countType = typeFactory.createSqlType(SqlTypeName.BIGINT);
+        RelDataType nullableInput = typeFactory.createTypeWithNullability(inputType, true);
+
+        RelDataType sumType = ReturnTypes.AGG_SUM.inferReturnType(opBinding);
+        if (sumType == null) {
+            sumType = nullableInput;
+        } else {
+            sumType = typeFactory.createTypeWithNullability(sumType, true);
+        }
+
+        RelDataType doubleNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.DOUBLE), true);
+
+        // Nested std_deviation_bounds struct.
+        RelDataType boundsType = typeFactory.createTypeWithNullability(
+            typeFactory.builder()
+                .add(EXTENDED_STATS_BOUND_UPPER, doubleNullable)
+                .add(EXTENDED_STATS_BOUND_LOWER, doubleNullable)
+                .add(EXTENDED_STATS_BOUND_UPPER_POPULATION, doubleNullable)
+                .add(EXTENDED_STATS_BOUND_LOWER_POPULATION, doubleNullable)
+                .add(EXTENDED_STATS_BOUND_UPPER_SAMPLING, doubleNullable)
+                .add(EXTENDED_STATS_BOUND_LOWER_SAMPLING, doubleNullable)
+                .build(),
+            true
+        );
+
+        return typeFactory.builder()
+            .add(STATS_FIELD_COUNT, countType)
+            .add(STATS_FIELD_MIN, nullableInput)
+            .add(STATS_FIELD_MAX, nullableInput)
+            .add(STATS_FIELD_AVG, doubleNullable)
+            .add(STATS_FIELD_SUM, sumType)
+            .add(EXTENDED_STATS_FIELD_SUM_OF_SQUARES, sumType)
+            .add(EXTENDED_STATS_FIELD_VARIANCE, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_VARIANCE_POPULATION, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_VARIANCE_SAMPLING, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_STD_DEVIATION, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_STD_DEVIATION_POPULATION, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_STD_DEVIATION_SAMPLING, doubleNullable)
+            .add(EXTENDED_STATS_FIELD_STD_DEVIATION_BOUNDS, boundsType)
+            .build();
+    }
 }

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

@@ -29,6 +29,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateReduceRule;
+import org.opensearch.analytics.planner.rules.OpenSearchExtendedStatsReduceRule;
 import org.opensearch.analytics.planner.rules.OpenSearchStatsReduceRule;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateRule;
 import org.opensearch.analytics.planner.rules.OpenSearchAggregateSplitRule;
@@ -187,6 +188,7 @@ private static RelNode decomposeAggregates(RelNode input, RuleProfilingListener
         // so the AVG decomposition rule does not see — and does not need to see — STATS' avg
         // path. Other AVG/STDDEV/VAR aggCalls in the same Aggregate are still picked up.
         builder.addRuleInstance(OpenSearchStatsReduceRule.INSTANCE);
+        builder.addRuleInstance(OpenSearchExtendedStatsReduceRule.INSTANCE);
         builder.addRuleInstance(new OpenSearchAggregateReduceRule());
         HepPlanner planner = new HepPlanner(builder.build());
         if (listener != null) {