grouping functions are now allowed as arguments to scalar functions

momovdg · momovdg · commit b582db0350c7 · 2026-01-30T10:56:42.000+02:00
diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/registry/stats/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/registry/stats/autocomplete.test.ts
@@ -256,7 +256,7 @@ describe('STATS Autocomplete', () => {
           [
             ...getFieldNamesByType([...ESQL_COMMON_NUMERIC_TYPES, 'date', 'date_nanos']),
             ...getFunctionSignaturesByReturnType(
-              Location.EVAL,
+              Location.STATS_BY,
               ['date', 'date_nanos', ...ESQL_COMMON_NUMERIC_TYPES],
               {
                 scalar: true,
@@ -277,9 +277,9 @@ describe('STATS Autocomplete', () => {
             }),
             ...getFieldNamesByType(roundParameterTypes),
             ...getFunctionSignaturesByReturnType(
-              Location.EVAL,
+              Location.STATS_BY,
               roundParameterTypes,
-              { scalar: true },
+              { scalar: true, grouping: true },
               undefined,
               ['round']
             ),
@@ -294,9 +294,9 @@ describe('STATS Autocomplete', () => {
             }),
             ...getFieldNamesByType(roundParameterTypes),
             ...getFunctionSignaturesByReturnType(
-              Location.EVAL,
+              Location.STATS_BY,
               ESQL_NUMBER_TYPES,
-              { scalar: true },
+              { scalar: true, grouping: true },
               undefined,
               ['round']
             ),
@@ -308,9 +308,9 @@ describe('STATS Autocomplete', () => {
           [
             ...getFieldNamesByType(roundParameterTypes),
             ...getFunctionSignaturesByReturnType(
-              Location.STATS,
+              Location.STATS_BY,
               ESQL_NUMBER_TYPES,
-              { scalar: true },
+              { scalar: true, grouping: true },
               undefined,
               ['round']
             ),
@@ -325,8 +325,9 @@ describe('STATS Autocomplete', () => {
           'from a | stats avg(',
           [
             ...expectedFieldsAvg,
-            ...getFunctionSignaturesByReturnType(Location.STATS, AVG_TYPES, {
+            ...getFunctionSignaturesByReturnType(Location.STATS_BY, AVG_TYPES, {
               scalar: true,
+              grouping: true,
             }),
           ],
           mockCallbacks
@@ -352,9 +353,9 @@ describe('STATS Autocomplete', () => {
           [
             ...expectedFieldsAvg,
             ...getFunctionSignaturesByReturnType(
-              Location.EVAL,
+              Location.STATS_BY,
               AVG_TYPES,
-              { scalar: true },
+              { scalar: true, grouping: true },
               undefined,
               ['round']
             ),
@@ -395,6 +396,7 @@ describe('STATS Autocomplete', () => {
             ],
             {
               scalar: true,
+              grouping: true,
             }
           ),
         ];
@@ -414,10 +416,11 @@ describe('STATS Autocomplete', () => {
           [
             ...getFieldNamesByType(AVG_TYPES),
             ...getFunctionSignaturesByReturnType(
-              Location.EVAL,
+              Location.STATS_BY,
               [...AVG_TYPES, 'aggregate_metric_double'],
               {
                 scalar: true,
+                grouping: true,
               }
             ),
           ],
@@ -570,11 +573,6 @@ describe('STATS Autocomplete', () => {
         await statsExpectSuggestions('from a | stats a=min(b) by ', expected);
       });
 
-      test('no grouping functions as args to scalar function', async () => {
-        const suggestions = await suggest('FROM a | STATS a=MIN(b) BY ACOS(');
-        expect(suggestions.some((s) => allGroupingFunctions.includes(s.text))).toBe(false);
-      });
-
       test('on partial column name', async () => {
         const expected = [
           ...allEvalFunctionsForStats,
diff --git a/src/platform/packages/shared/kbn-esql-language/src/commands/registry/stats/autocomplete.ts b/src/platform/packages/shared/kbn-esql-language/src/commands/registry/stats/autocomplete.ts
@@ -80,7 +80,7 @@ export async function autocomplete(
       let location: Location;
 
       if (isInBy) {
-        location = Location.EVAL;
+        location = Location.STATS_BY;
       } else if (isTimeseriesSource && isAggFunctionUsedAlready(command, command.args.length - 1)) {
         location = Location.STATS_TIMESERIES;
       } else {
@@ -400,9 +400,11 @@ function buildCustomFilteringContext(
   }
 
   const statsSpecificFunctionsToIgnore: string[] = [];
-  statsSpecificFunctionsToIgnore.push(
-    ...getAllFunctions({ type: FunctionDefinitionTypes.GROUPING }).map(({ name }) => name)
-  );
+  if (basicContext.functionDefinition?.type === 'grouping') {
+    statsSpecificFunctionsToIgnore.push(
+      ...getAllFunctions({ type: FunctionDefinitionTypes.GROUPING }).map(({ name }) => name)
+    );
+  }
 
   const finalCommandArgIndex = command.args.length - 1;
   const isInBy = isNodeWithinByClause(foundFunction, command);
