PPL - Fix count/dc(field) failing after a fields command

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -1616,7 +1616,18 @@ private Pair<List<RexNode>, List<AggCall>> 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()));
+        // This filter is stacked on top of the Project, so its reference has to address the
+        // Project's OUTPUT. distinctRefsOfCounts may instead hold an index mapped through the
+        // Project, which addresses its INPUT — that mapping exists only to prove two aliases are
+        // one column, and it is not a reference we can use here. `fields COMM | stats count(COMM)`
+        // is the clearest case: COMM is $6 of the scan and $0 of the projection, and filtering on
+        // $6 fails with "RexInputRef index 6 out of range 0..0". Where the projection merely
+        // reorders, the index stays in range and names a different column, so the wrong column's
+        // nullness is tested. refsPerCount is already in the output frame, and by this branch every
+        // entry denotes the same column, so any one of them gives the same predicate.
+        RexInputRef filterRef =
+            refsPerCount.stream().flatMap(List::stream).findFirst().orElseThrow();
+        context.relBuilder.filter(context.relBuilder.isNotNull(filterRef));
       }
     }
 

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)])

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLCountFrameTest.java

@@ -0,0 +1,76 @@
+/*
+ * 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;
+
+/**
+ * {@code stats count(field)} adds an is-not-null filter, and these pin the frame its reference is
+ * expressed in.
+ *
+ * <p>The filter is stacked on whatever the builder is holding, so the reference has to address that
+ * node's output. A {@code fields} command in front of the {@code stats} makes the difference
+ * visible: the column's index in the projection's output is not its index in the projection's
+ * input, and using the latter is out of range as soon as the projection is narrower than its input.
+ */
+public class CalcitePPLCountFrameTest extends CalcitePPLAbstractTest {
+
+  public CalcitePPLCountFrameTest() {
+    super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL);
+  }
+
+  /** COMM is $6 of the scan but $0 of the projection the filter sits on. */
+  @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");
+  }
+
+  /** Two columns kept, so the index is in range either way but only one of them is COMM. */
+  @Test
+  public void testCountAfterReorderingFields() {
+    String ppl = "source=EMP | fields SAL, COMM | stats count(COMM) as c";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        ""
+            + "LogicalAggregate(group=[{}], c=[COUNT($0)])\n"
+            + "  LogicalProject(COMM=[$1])\n"
+            + "    LogicalFilter(condition=[IS NOT NULL($1)])\n"
+            + "      LogicalProject(SAL=[$5], COMM=[$6])\n"
+            + "        LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+    verifyResult(root, "c=4\n");
+  }
+
+  /** dc() takes the same filter, so it takes the same 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 the case the index mapping exists for: it recognises them as
+   * the same column so a single filter covers both counts.
+   */
+  @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");
+  }
+}