diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/AggregationOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/AggregationOperator.java index 5e05286bbce..d71089d9901 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/AggregationOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/AggregationOperator.java @@ -55,14 +55,19 @@ public AggregationOperator(PhysicalPlan input, List aggregatorL List groupByExprList) { this.input = input; this.aggregatorList = aggregatorList; + this.groupByExprList = groupByExprList; if (hasSpan(groupByExprList)) { + // span expression is always the first expression in group list if exist. this.span = groupByExprList.get(0); - this.groupByExprList = groupByExprList.subList(1, groupByExprList.size()); + this.collector = + Collector.Builder.build( + this.span, groupByExprList.subList(1, groupByExprList.size()), this.aggregatorList); + } else { this.span = null; - this.groupByExprList = groupByExprList; + this.collector = + Collector.Builder.build(this.span, this.groupByExprList, this.aggregatorList); } - this.collector = Collector.Builder.build(this.span, this.groupByExprList, this.aggregatorList); } @Override diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/AggregationOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/AggregationOperatorTest.java index 3b45a11c6ce..318499c0759 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/AggregationOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/AggregationOperatorTest.java @@ -495,4 +495,17 @@ public void twoBucketsSpanAndLong() { "span", new ExprDateValue("2021-01-07"), "region","iad", "host", "h2", "max", 8)) )); } + + @Test + public void copyOfAggregationOperatorShouldSame() { + AggregationOperator plan = new AggregationOperator(testScan(datetimeInputs), + Collections.singletonList(DSL + .named("count", dsl.count(DSL.ref("second", TIMESTAMP)))), + Collections.singletonList(DSL + .named("span", DSL.span(DSL.ref("second", TIMESTAMP), DSL.literal(6 * 1000), "ms")))); + AggregationOperator copy = new AggregationOperator(plan.getInput(), plan.getAggregatorList(), + plan.getGroupByExprList()); + + assertEquals(plan, copy); + } }