Why decompose at HEP time? Same as {@code OpenSearchStatsReduceRule}: + * {@code EXTENDED_STATS} has no executor implementation in this codebase and isn't a + * Calcite standard {@link org.apache.calcite.sql.SqlKind}. Doing the rewrite during HEP + * keeps the rest of the planner unaware of the bundle aggregate; everything downstream + * sees only primitives that already have full support. + * + * @opensearch.internal + */ +public final class OpenSearchExtendedStatsReduceRule extends RelOptRule { + + public static final OpenSearchExtendedStatsReduceRule INSTANCE = new OpenSearchExtendedStatsReduceRule(); + + /** Synthetic name prefix for the squared-value columns added beneath the rebuilt aggregate. */ + private static final String SQUARED_COLUMN_NAME_PREFIX = "$ext_stats_sq_"; + + private OpenSearchExtendedStatsReduceRule() { + super(operand(LogicalAggregate.class, any()), "OpenSearchExtendedStatsReduceRule"); + } + + @Override + public boolean matches(RelOptRuleCall call) { + LogicalAggregate aggregate = call.rel(0); + for (AggregateCall aggCall : aggregate.getAggCallList()) { + if (isExtendedStats(aggCall.getAggregation())) { + return true; + } + } + return false; + } + + @Override + public void onMatch(RelOptRuleCall call) { + LogicalAggregate aggregate = call.rel(0); + RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder(); + RelDataTypeFactory typeFactory = aggregate.getCluster().getTypeFactory(); + + int groupCount = aggregate.getGroupCount(); + boolean hasEmptyGroup = aggregate.getGroupSet().isEmpty(); + + Expansion expansion = expandAggCalls(aggregate, hasEmptyGroup, groupCount, rexBuilder); + + LogicalAggregate newAgg = LogicalAggregate.create( + expansion.aggInput, + aggregate.getHints(), + aggregate.getGroupSet(), + aggregate.getGroupSets(), + expansion.newCalls + ); + + RelNode project = buildAssemblyProject(aggregate, newAgg, expansion.originalToNew, groupCount, rexBuilder, typeFactory); + call.transformTo(project); + } + + /** + * Expands each EXTENDED_STATS aggregate call into five primitives ({@code SUM(x)}, + * {@code MIN(x)}, {@code MAX(x)}, {@code SUM(x_squared)}, {@code COUNT(x)}, in that + * emission order); leaves non-EXTENDED_STATS calls untouched. + * + *
The {@code x_squared} input is computed by a {@link LogicalProject} inserted + * beneath the rebuilt aggregate, which appends one squared-value column per + * unique EXTENDED_STATS source column. The Project is shared by every aggregate that + * needs that derivation. Putting the multiplication into a Project beneath rather than + * into the aggregate call's {@code rexList} matters because + * {@code DistributedAggregateRewriter} preserves {@code rexList} verbatim into the + * FINAL aggregate after Volcano split — and at FINAL time the rexList's input refs + * would point at the wrong columns of the gathered partial output, producing + * semantically incorrect results. With the Project beneath, FINAL's argList rebases + * naturally to the column where PARTIAL's {@code sum_of_squares} output lands. + * + *
Same SUM-first emission rationale as {@code OpenSearchStatsReduceRule} — keeps
+ * {@code typeMatchesInferred} happy after Volcano split when FINAL re-infers types
+ * against PARTIAL's first column.
+ */
+ private static Expansion expandAggCalls(
+ LogicalAggregate aggregate,
+ boolean hasEmptyGroup,
+ int groupCount,
+ RexBuilder rexBuilder
+ ) {
+ RelNode input = aggregate.getInput();
+
+ Map Variance and std_deviation triplets duplicate the population variant because legacy
+ * defines {@code variance} as an alias for {@code variance_population} (same value), and
+ * likewise for std_deviation. The sampling variants use Bessel's correction ({@code n-1}
+ * denominator) and return NULL when {@code count <= 1}.
+ */
+ private static RexNode buildExtendedStatsStruct(
+ AggregateCall original,
+ int[] mapping,
+ List By the time these assertions run, {@code OpenSearchExtendedStatsReduceRule} has
+ * expanded EXTENDED_STATS into the five primitive aggregate calls plus a Project that
+ * builds the 13-field struct (count, min, max, avg, sum, sum_of_squares, variance triplet,
+ * std_deviation triplet, std_deviation_bounds nested struct). Subsequent phases operate
+ * on that primitive plan exactly the way they operate on any other SUM/COUNT aggregate.
+ */
+public class ExtendedStatsAggregatePlanShapeTests extends PlanShapeTestBase {
+
+ public void testExtendedStatsSingleField_1shard_yieldsProjectOverSingleAggregate() {
+ RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+ AggregateCall ext = makeExtendedStatsCall(scan, 1, "ext_size");
+ RelNode plan = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(ext));
+
+ RelNode result = runPlanner(plan, singleShardContext());
+
+ // Top-level: ExchangeReducer wrapping a Project that builds the struct from a
+ // single-mode Aggregate's primitives. Peel the ER to inspect Project + Aggregate.
+ RelNode top = unwrapRootReducer(result);
+ assertTrue("Top must be OpenSearchProject", top instanceof OpenSearchProject);
+ OpenSearchProject project = (OpenSearchProject) top;
+ assertEquals(1, project.getProjects().size());
+
+ RelNode child = unwrapExchange(project.getInput());
+ assertTrue("Project input must be OpenSearchAggregate", child instanceof OpenSearchAggregate);
+ OpenSearchAggregate agg = (OpenSearchAggregate) child;
+ assertEquals("Decomposed to 5 primitive calls", 5, agg.getAggCallList().size());
+ }
+
+ public void testExtendedStatsSingleField_2shard_splitsIntoPartialFinalAroundProject() {
+ RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+ AggregateCall ext = makeExtendedStatsCall(scan, 1, "ext_size");
+ RelNode plan = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(ext));
+
+ RelNode result = runPlanner(plan, multiShardContext());
+
+ // Plan shape after the multi-shard split:
+ // Project [ext = STRUCT(count, min, max, avg, sum, sum_of_squares,
+ // variance, variance_population, variance_sampling,
+ // std_deviation, std_deviation_population, std_deviation_sampling,
+ // std_deviation_bounds)]
+ // └── Aggregate(FINAL) [SUM, MIN, MAX, SUM, COUNT]
+ // └── ExchangeReducer
+ // └── Aggregate(PARTIAL) [SUM, MIN, MAX, SUM, COUNT]
+ // └── Project [orig cols + x*x] ← derives sum_of_squares input
+ // └── Scan
+ RelNode top = unwrapRootReducer(result);
+ assertTrue("Top must be OpenSearchProject", top instanceof OpenSearchProject);
+ OpenSearchProject project = (OpenSearchProject) top;
+
+ RelNode finalAgg = unwrapExchange(project.getInput());
+ assertTrue("Below project must be Aggregate(FINAL)", finalAgg instanceof OpenSearchAggregate);
+ assertEquals(5, ((OpenSearchAggregate) finalAgg).getAggCallList().size());
+
+ RelNode reducer = finalAgg.getInputs().get(0);
+ assertTrue("Below FINAL must be ExchangeReducer", reducer instanceof OpenSearchExchangeReducer);
+ }
+
+ public void testExtendedStatsWithGroupBy_2shard() {
+ RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+ AggregateCall ext = makeExtendedStatsCall(scan, 1, "ext_size");
+ RelNode plan = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(ext));
+
+ RelNode result = runPlanner(plan, multiShardContext());
+
+ RelNode top = unwrapRootReducer(result);
+ assertTrue("Top must be OpenSearchProject", top instanceof OpenSearchProject);
+ OpenSearchProject project = (OpenSearchProject) top;
+ assertEquals(2, project.getProjects().size());
+ assertEquals("status", project.getRowType().getFieldList().get(0).getName());
+ assertEquals("ext_size", project.getRowType().getFieldList().get(1).getName());
+ }
+
+ public void testExtendedStatsMixedWithCount_2shard_bothFlowThroughSamePartialFinal() {
+ RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+ AggregateCall ext = makeExtendedStatsCall(scan, 1, "ext_size");
+ AggregateCall countStar = AggregateCall.create(
+ org.apache.calcite.sql.fun.SqlStdOperatorTable.COUNT,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ scan,
+ null,
+ "cnt_all"
+ );
+ RelNode plan = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(ext, countStar));
+
+ RelNode result = runPlanner(plan, multiShardContext());
+
+ RelNode top = unwrapRootReducer(result);
+ assertTrue("Top must be OpenSearchProject", top instanceof OpenSearchProject);
+ OpenSearchProject project = (OpenSearchProject) top;
+ assertEquals(3, project.getProjects().size());
+
+ RelNode finalAgg = unwrapExchange(project.getInput());
+ assertTrue("Below project must be Aggregate(FINAL)", finalAgg instanceof OpenSearchAggregate);
+ // 5 primitives for EXTENDED_STATS + 1 for cnt_all = 6.
+ assertEquals(6, ((OpenSearchAggregate) finalAgg).getAggCallList().size());
+ }
+
+ // ---- Helpers ----
+
+ private AggregateCall makeExtendedStatsCall(RelNode input, int colIdx, String name) {
+ return AggregateCall.create(
+ OpenSearchAggregateOperators.EXTENDED_STATS,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(colIdx),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ input,
+ null,
+ name
+ );
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchExtendedStatsReduceRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchExtendedStatsReduceRuleTests.java
new file mode 100644
index 0000000000000..cbc309ef2db20
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchExtendedStatsReduceRuleTests.java
@@ -0,0 +1,379 @@
+/*
+ * 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.planner.rules;
+
+import org.apache.calcite.plan.hep.HepMatchOrder;
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.core.Project;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.opensearch.analytics.planner.BasePlannerRulesTests;
+import org.opensearch.analytics.spi.OpenSearchAggregateOperators;
+
+import java.util.List;
+
+/**
+ * Unit tests for {@link OpenSearchExtendedStatsReduceRule}. Mirrors
+ * {@code OpenSearchStatsReduceRuleTests} — applies only this rule via a minimal
+ * {@link HepPlanner} and asserts the resulting tree's structure.
+ *
+ * Tests are intentionally rule-scoped — they do not exercise {@code PlannerImpl}
+ * marking, CBO, or fragment conversion. End-to-end planner integration is covered by
+ * {@code ExtendedStatsAggregatePlanShapeTests}.
+ */
+public class OpenSearchExtendedStatsReduceRuleTests extends BasePlannerRulesTests {
+
+ public void testSingleExtendedStatsCall_noGroupBy_decomposesIntoFivePrimitivesPlusProject() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall call = makeExtendedStatsCall(scan, 0, "ext_stats_price");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ assertProjectOverAggregateWithFivePrimitives(rewritten, 0);
+ }
+
+ public void testSingleExtendedStatsCall_withGroupBy_preservesGroupSetAndPassesThroughGroupColumn() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "city", "price"));
+ AggregateCall call = makeExtendedStatsCall(scan, 1, "ext_stats_price");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ assertTrue("Top must be a Project", rewritten instanceof Project);
+ Project project = (Project) rewritten;
+ assertEquals(2, project.getProjects().size());
+ assertEquals("city", project.getRowType().getFieldList().get(0).getName());
+ assertEquals("ext_stats_price", project.getRowType().getFieldList().get(1).getName());
+
+ Aggregate inner = (Aggregate) project.getInput();
+ assertEquals(ImmutableBitSet.of(0), inner.getGroupSet());
+ // Order matters — SUM-first emission preserves typeMatchesInferred after Volcano split.
+ assertEquals(5, inner.getAggCallList().size());
+ assertAggCallKinds(inner, SqlKind.SUM, SqlKind.MIN, SqlKind.MAX, SqlKind.SUM, SqlKind.COUNT);
+ }
+
+ public void testExtendedStatsMixedWithOtherAggregate_decomposesOnlyExtendedStats() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price", "quantity"));
+ AggregateCall extStatsCall = makeExtendedStatsCall(scan, 0, "ext_stats_price");
+ AggregateCall avgCall = AggregateCall.create(
+ SqlStdOperatorTable.AVG,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(1),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ scan,
+ null,
+ "avg_qty"
+ );
+ LogicalAggregate aggregate = LogicalAggregate.create(
+ scan,
+ List.of(),
+ ImmutableBitSet.of(),
+ null,
+ List.of(extStatsCall, avgCall)
+ );
+
+ RelNode rewritten = applyRule(aggregate);
+
+ Project project = (Project) rewritten;
+ assertEquals(2, project.getProjects().size());
+
+ Aggregate inner = (Aggregate) project.getInput();
+ // 5 primitives for EXTENDED_STATS + 1 surviving AVG = 6.
+ assertEquals(6, inner.getAggCallList().size());
+ assertAggCallKinds(inner, SqlKind.SUM, SqlKind.MIN, SqlKind.MAX, SqlKind.SUM, SqlKind.COUNT, SqlKind.AVG);
+ assertSame(SqlStdOperatorTable.AVG, inner.getAggCallList().get(5).getAggregation());
+ }
+
+ public void testMultipleExtendedStatsCalls_eachExpandsIntoFivePrimitives() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price", "quantity"));
+ AggregateCall ext1 = makeExtendedStatsCall(scan, 0, "ext_price");
+ AggregateCall ext2 = makeExtendedStatsCall(scan, 1, "ext_qty");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(ext1, ext2));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ Project project = (Project) rewritten;
+ assertEquals(2, project.getProjects().size());
+
+ Aggregate inner = (Aggregate) project.getInput();
+ // 5 primitives per EXTENDED_STATS × 2 = 10.
+ assertEquals(10, inner.getAggCallList().size());
+ assertAggCallKinds(
+ inner,
+ SqlKind.SUM,
+ SqlKind.MIN,
+ SqlKind.MAX,
+ SqlKind.SUM,
+ SqlKind.COUNT,
+ SqlKind.SUM,
+ SqlKind.MIN,
+ SqlKind.MAX,
+ SqlKind.SUM,
+ SqlKind.COUNT
+ );
+ assertTrue(project.getProjects().get(0) instanceof RexCall);
+ assertEquals(SqlKind.ROW, ((RexCall) project.getProjects().get(0)).getOperator().getKind());
+ assertTrue(project.getProjects().get(1) instanceof RexCall);
+ assertEquals(SqlKind.ROW, ((RexCall) project.getProjects().get(1)).getOperator().getKind());
+ }
+
+ public void testSumOfSquaresAggregatesOverDerivedSquaredColumn() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall call = makeExtendedStatsCall(scan, 0, "ext_stats_price");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ // The rebuilt aggregate's input must be a LogicalProject that exposes the original
+ // columns plus one derived squared column. Putting x*x in a Project beneath (not in
+ // the aggCall's rexList) is essential: DistributedAggregateRewriter preserves rexList
+ // verbatim into FINAL after Volcano split, which would silently change the meaning of
+ // SUM(x*x). Routing through a derived column lets FINAL's argList rebase normally.
+ Aggregate inner = (Aggregate) ((Project) rewritten).getInput();
+ assertTrue(
+ "Aggregate's input must be a LogicalProject that adds the squared column",
+ inner.getInput() instanceof LogicalProject
+ );
+ LogicalProject squaredProject = (LogicalProject) inner.getInput();
+ assertEquals("Project must add 1 squared column to the original 1-column input", 2, squaredProject.getProjects().size());
+
+ RexNode squaredExpr = squaredProject.getProjects().get(1);
+ assertTrue(squaredExpr instanceof RexCall);
+ assertSame(SqlStdOperatorTable.MULTIPLY, ((RexCall) squaredExpr).getOperator());
+
+ AggregateCall sumOfSquares = inner.getAggCallList().get(3);
+ assertSame(SqlStdOperatorTable.SUM, sumOfSquares.getAggregation());
+ assertTrue("SUM_OF_SQUARES must have empty rexList", sumOfSquares.rexList.isEmpty());
+ assertEquals("SUM_OF_SQUARES argList must reference the squared column", List.of(1), sumOfSquares.getArgList());
+ }
+
+ public void testMultipleExtendedStatsOnSameSourceShareOneSquaredColumn() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall ext1 = makeExtendedStatsCall(scan, 0, "ext_a");
+ AggregateCall ext2 = makeExtendedStatsCall(scan, 0, "ext_b");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(ext1, ext2));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ Aggregate inner = (Aggregate) ((Project) rewritten).getInput();
+ LogicalProject squaredProject = (LogicalProject) inner.getInput();
+ assertEquals("Same-source EXTENDED_STATS calls must share one squared column", 2, squaredProject.getProjects().size());
+ assertEquals(List.of(1), inner.getAggCallList().get(3).getArgList());
+ assertEquals(List.of(1), inner.getAggCallList().get(8).getArgList());
+ }
+
+ public void testExtendedStatsCall_filterArgIsPropagatedToAllPrimitives() {
+ // mockTable's default schema is single-column INTEGER, but FILTER needs a BOOLEAN
+ // NOT NULL column — build a 2-col table where col 1 is BOOLEAN.
+ StubTableScan scan = (StubTableScan) stubScan(
+ mockTable("orders2", new String[] { "price", "in_scope" }, new SqlTypeName[] { SqlTypeName.INTEGER, SqlTypeName.BOOLEAN })
+ );
+ AggregateCall call = AggregateCall.create(
+ OpenSearchAggregateOperators.EXTENDED_STATS,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(0),
+ 1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ scan,
+ null,
+ "ext_stats_price"
+ );
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ Aggregate inner = (Aggregate) ((Project) rewritten).getInput();
+ for (AggregateCall primitive : inner.getAggCallList()) {
+ assertEquals("filterArg must propagate to " + primitive.getAggregation().getName(), 1, primitive.filterArg);
+ }
+ }
+
+ public void testRuleNoOpWhenNoExtendedStatsCallPresent() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall sumCall = AggregateCall.create(
+ SqlStdOperatorTable.SUM,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(0),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ scan,
+ null,
+ "sum_price"
+ );
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sumCall));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ // No EXTENDED_STATS → rule does not fire → HEP returns the input unchanged.
+ assertTrue("Plan must remain a LogicalAggregate", rewritten instanceof LogicalAggregate);
+ assertEquals(1, ((LogicalAggregate) rewritten).getAggCallList().size());
+ assertSame(SqlStdOperatorTable.SUM, ((LogicalAggregate) rewritten).getAggCallList().get(0).getAggregation());
+ }
+
+ public void testExtendedStatsStructFieldNamesMatchSpec() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall call = makeExtendedStatsCall(scan, 0, "ext_stats_price");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ // Order matches legacy OpenSearch InternalExtendedStats.toXContent.
+ Project project = (Project) rewritten;
+ RelDataType type = project.getRowType().getFieldList().get(0).getType();
+ assertTrue(type.isStruct());
+ assertEquals(13, type.getFieldList().size());
+ assertEquals(OpenSearchAggregateOperators.STATS_FIELD_COUNT, type.getFieldList().get(0).getName());
+ assertEquals(OpenSearchAggregateOperators.STATS_FIELD_MIN, type.getFieldList().get(1).getName());
+ assertEquals(OpenSearchAggregateOperators.STATS_FIELD_MAX, type.getFieldList().get(2).getName());
+ assertEquals(OpenSearchAggregateOperators.STATS_FIELD_AVG, type.getFieldList().get(3).getName());
+ assertEquals(OpenSearchAggregateOperators.STATS_FIELD_SUM, type.getFieldList().get(4).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_SUM_OF_SQUARES, type.getFieldList().get(5).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_VARIANCE, type.getFieldList().get(6).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_VARIANCE_POPULATION, type.getFieldList().get(7).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_VARIANCE_SAMPLING, type.getFieldList().get(8).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_STD_DEVIATION, type.getFieldList().get(9).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_STD_DEVIATION_POPULATION, type.getFieldList().get(10).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_STD_DEVIATION_SAMPLING, type.getFieldList().get(11).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_FIELD_STD_DEVIATION_BOUNDS, type.getFieldList().get(12).getName());
+ }
+
+ public void testExtendedStatsBoundsStructFieldNamesMatchSpec() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall call = makeExtendedStatsCall(scan, 0, "ext_stats_price");
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ // Order matches the std_deviation_bounds sub-struct of legacy InternalExtendedStats.
+ Project project = (Project) rewritten;
+ RelDataType extType = project.getRowType().getFieldList().get(0).getType();
+ RelDataType boundsType = extType.getFieldList().get(12).getType();
+ assertTrue("std_deviation_bounds must be a struct", boundsType.isStruct());
+ assertEquals(6, boundsType.getFieldList().size());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_UPPER, boundsType.getFieldList().get(0).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_LOWER, boundsType.getFieldList().get(1).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_UPPER_POPULATION, boundsType.getFieldList().get(2).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_LOWER_POPULATION, boundsType.getFieldList().get(3).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_UPPER_SAMPLING, boundsType.getFieldList().get(4).getName());
+ assertEquals(OpenSearchAggregateOperators.EXTENDED_STATS_BOUND_LOWER_SAMPLING, boundsType.getFieldList().get(5).getName());
+ }
+
+ public void testExtendedStatsCall_distinctIsPropagatedWhereMeaningful() {
+ StubTableScan scan = (StubTableScan) stubScan(mockTable("orders", "price"));
+ AggregateCall call = AggregateCall.create(
+ OpenSearchAggregateOperators.EXTENDED_STATS,
+ true,
+ false,
+ false,
+ List.of(),
+ List.of(0),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ scan,
+ null,
+ "ext_stats_price"
+ );
+ LogicalAggregate aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(call));
+
+ RelNode rewritten = applyRule(aggregate);
+
+ Aggregate inner = (Aggregate) ((Project) rewritten).getInput();
+ // Calcite normalizes DISTINCT off MIN/MAX (idempotent). SUM and COUNT preserve it.
+ for (AggregateCall primitive : inner.getAggCallList()) {
+ SqlKind kind = primitive.getAggregation().getKind();
+ if (kind == SqlKind.SUM || kind == SqlKind.COUNT) {
+ assertTrue("DISTINCT must propagate to " + kind, primitive.isDistinct());
+ } else {
+ assertFalse("Calcite normalizes DISTINCT off " + kind, primitive.isDistinct());
+ }
+ }
+ }
+
+ // ---- Helpers ----
+
+ private AggregateCall makeExtendedStatsCall(RelNode input, int colIdx, String name) {
+ return AggregateCall.create(
+ OpenSearchAggregateOperators.EXTENDED_STATS,
+ false,
+ false,
+ false,
+ List.of(),
+ List.of(colIdx),
+ -1,
+ null,
+ RelCollations.EMPTY,
+ 0,
+ input,
+ null,
+ name
+ );
+ }
+
+ private RelNode applyRule(RelNode root) {
+ HepProgramBuilder builder = new HepProgramBuilder();
+ builder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ builder.addRuleInstance(OpenSearchExtendedStatsReduceRule.INSTANCE);
+ HepPlanner planner = new HepPlanner(builder.build());
+ planner.setRoot(root);
+ return planner.findBestExp();
+ }
+
+ private static void assertProjectOverAggregateWithFivePrimitives(RelNode rewritten, int colIdx) {
+ assertTrue("Top must be a Project", rewritten instanceof Project);
+ Project project = (Project) rewritten;
+ assertTrue("Project's input must be a LogicalAggregate", project.getInput() instanceof LogicalAggregate);
+ LogicalAggregate inner = (LogicalAggregate) project.getInput();
+ assertEquals(5, inner.getAggCallList().size());
+ assertAggCallKinds(inner, SqlKind.SUM, SqlKind.MIN, SqlKind.MAX, SqlKind.SUM, SqlKind.COUNT);
+ RelDataType extType = project.getRowType().getFieldList().get(colIdx).getType();
+ assertTrue("EXTENDED_STATS column must be a struct type", extType.isStruct());
+ }
+
+ private static void assertAggCallKinds(Aggregate aggregate, SqlKind... expectedKinds) {
+ assertEquals("aggCall count mismatch", expectedKinds.length, aggregate.getAggCallList().size());
+ for (int i = 0; i < expectedKinds.length; i++) {
+ assertEquals(
+ "aggCall " + i + " kind mismatch",
+ expectedKinds[i],
+ aggregate.getAggCallList().get(i).getAggregation().getKind()
+ );
+ }
+ }
+}