diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index a2f0addbfb6..e66e3319b57 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -1616,7 +1616,13 @@ private Pair, List> aggregateWithTrimming( distinctRefsOfCounts = refsPerCount.stream().flatMap(List::stream).distinct().toList(); } if (distinctRefsOfCounts.size() == 1 && refsPerCount.stream().noneMatch(List::isEmpty)) { - context.relBuilder.filter(context.relBuilder.isNotNull(distinctRefsOfCounts.getFirst())); + // The filter is stacked on the Project, so its reference must address the Project's output. + // distinctRefsOfCounts may hold an index mapped through the Project, which addresses its + // input; that mapping only serves to prove two aliases are one column. refsPerCount is + // already in the output frame, and every entry here denotes the same column. + RexInputRef filterRef = + refsPerCount.stream().flatMap(List::stream).findFirst().orElseThrow(); + context.relBuilder.filter(context.relBuilder.isNotNull(filterRef)); } } diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_count_agg_push3.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_count_agg_push3.yaml index bd1fff2449a..0e119e2b819 100644 --- a/integ-test/src/test/resources/expectedOutput/calcite/explain_count_agg_push3.yaml +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_count_agg_push3.yaml @@ -3,8 +3,8 @@ calcite: LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT]) LogicalAggregate(group=[{}], cnt=[COUNT($0)]) LogicalProject(name=[$17]) - LogicalFilter(condition=[IS NOT NULL($10)]) + LogicalFilter(condition=[IS NOT NULL($17)]) LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], name=[$10]) CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]]) physical: | - CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={},cnt=COUNT($1)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"exists":{"field":"lastname","boost":1.0}},"track_total_hits":2147483647}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) + CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={},cnt=COUNT($0)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"exists":{"field":"lastname","boost":1.0}},"track_total_hits":2147483647}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)]) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLCountFrameTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLCountFrameTest.java new file mode 100644 index 00000000000..90621da5bbf --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLCountFrameTest.java @@ -0,0 +1,77 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.calcite; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.test.CalciteAssert; +import org.junit.Test; + +/** + * The is-not-null filter that {@code stats count(field)} adds sits on the projection the aggregate + * reads, so its reference must be the column's index in that projection's output. An input-side + * index is out of range when the projection narrows, and names the wrong column when it does not. + * + *

EMP is {@code EMPNO $0, ENAME $1, JOB $2, MGR $3, HIREDATE $4, SAL $5, COMM $6, DEPTNO $7}. + */ +public class CalcitePPLCountFrameTest extends CalcitePPLAbstractTest { + + public CalcitePPLCountFrameTest() { + super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL); + } + + /** COMM is $6 of the scan and $0 of the projection, so an input-side index is out of range. */ + @Test + public void testCountAfterNarrowingFields() { + String ppl = "source=EMP | fields COMM | stats count(COMM) as c"; + RelNode root = getRelNode(ppl); + String expectedLogical = + "" + + "LogicalAggregate(group=[{}], c=[COUNT($0)])\n" + + " LogicalFilter(condition=[IS NOT NULL($0)])\n" + + " LogicalProject(COMM=[$6])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + verifyResult(root, "c=4\n"); + } + + /** + * SAL sits at output position 6, so $6 stays in range and resolves to SAL while COMM is $5. Both + * are DECIMAL, so nothing objects. The result stays 4 because count ignores nulls itself, so only + * the plan shows which column the filter is about. + */ + @Test + public void testCountWhenTheStaleIndexStaysInRange() { + String ppl = + "source=EMP | fields EMPNO, ENAME, JOB, MGR, HIREDATE, COMM, SAL | stats count(COMM) as c"; + RelNode root = getRelNode(ppl); + String expectedLogical = + "" + + "LogicalAggregate(group=[{}], c=[COUNT($0)])\n" + + " LogicalProject(COMM=[$5])\n" + + " LogicalFilter(condition=[IS NOT NULL($5)])\n" + + " LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4]," + + " COMM=[$6], SAL=[$5])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + verifyLogical(root, expectedLogical); + verifyResult(root, "c=4\n"); + } + + /** dc() takes the same filter, so it took the same out-of-range reference. */ + @Test + public void testDistinctCountAfterNarrowingFields() { + String ppl = "source=EMP | fields COMM | stats dc(COMM) as c"; + RelNode root = getRelNode(ppl); + verifyResult(root, "c=4\n"); + } + + /** Two names for one column, which is what the input-side mapping is kept for. */ + @Test + public void testCountOfAnAliasedColumn() { + String ppl = "source=EMP | eval bonus = COMM | fields bonus | stats count(bonus) as c"; + RelNode root = getRelNode(ppl); + verifyResult(root, "c=4\n"); + } +}