ES|QL: manage INLINE STATS count(*) on result sets with no columns (#137017)

Author: luigidellaquila

docs/changelog/137017.yaml

@@ -0,0 +1,5 @@
+pr: 137017
+summary: Manage INLINE STATS count(*) on result sets with no columns
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/CountAggregatorFunction.java

@@ -78,23 +78,42 @@ public int intermediateBlockCount() {
     }
 
     private int blockIndex() {
-        return countAll ? 0 : channels.get(0);
+        // In case of countAll, block index is irrelevant.
+        // Page.positionCount should be used instead,
+        // because the page could have zero blocks
+        // (drop all columns scenario)
+        return countAll ? -1 : channels.get(0);
     }
 
     @Override
     public void addRawInput(Page page, BooleanVector mask) {
-        Block block = page.getBlock(blockIndex());
-        LongState state = this.state;
-        int count;
-        if (mask.isConstant()) {
-            if (mask.getBoolean(0) == false) {
-                return;
+        if (countAll) {
+            // this will work also when the page has no blocks
+            if (mask.isConstant() && mask.getBoolean(0)) {
+                state.longValue(state.longValue() + page.getPositionCount());
+            } else {
+                int count = 0;
+                for (int i = 0; i < mask.getPositionCount(); i++) {
+                    if (mask.getBoolean(i)) {
+                        count++;
+                    }
+                }
+                state.longValue(state.longValue() + count);
             }
-            count = countAll ? block.getPositionCount() : block.getTotalValueCount();
         } else {
-            count = countMasked(block, mask);
+            Block block = page.getBlock(blockIndex());
+            LongState state = this.state;
+            int count;
+            if (mask.isConstant()) {
+                if (mask.getBoolean(0) == false) {
+                    return;
+                }
+                count = block.getTotalValueCount();
+            } else {
+                count = countMasked(block, mask);
+            }
+            state.longValue(state.longValue() + count);
         }
-        state.longValue(state.longValue() + count);
     }
 
     private int countMasked(Block block, BooleanVector mask) {

x-pack/plugin/esql/qa/testFixtures/src/main/resources/inlinestats.csv-spec

@@ -4363,3 +4363,32 @@ from employees
 c:long
 1
 ;
+
+
+inlineStatsKeepDropCountStar
+required_capability: inline_stats_with_no_columns
+
+from employees
+| keep emp_no
+| drop emp_no
+| inline stats c = count(*)
+| limit 3
+;
+
+c:long
+100
+100
+100
+;
+
+rowInlineStatsKeepDropCountStar
+required_capability: inline_stats_with_no_columns
+
+row a = 1
+| drop a
+| inline stats c = count(*)
+;
+
+c:long
+1
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1569,11 +1569,17 @@ public enum Cap {
          * Temporarily forbid the use of an explicit or implicit LIMIT before INLINE STATS.
          */
         FORBID_LIMIT_BEFORE_INLINE_STATS(INLINE_STATS.enabled),
+
         /**
          * Support for the TRANGE function
          */
         FN_TRANGE,
 
+        /**
+         * https://github.com/elastic/elasticsearch/issues/136851
+         */
+        INLINE_STATS_WITH_NO_COLUMNS(INLINE_STATS.enabled),
+
         // Last capability should still have a comma for fewer merge conflicts when adding new ones :)
         // This comment prevents the semicolon from being on the previous capability when Spotless formats the file.
         ;