[ES|QL] Making the Expression Suggestor More Semantically Intelligent

src/platform/packages/shared/kbn-esql-ast/scripts/functions.ts

@@ -105,7 +105,10 @@ export function enrichFunctionSignatures(
  * - `date_extract`: Adds datePart suggestions from dateExtractOptions
  * - `mv_sort`: Marks 'order' as constantOnly with ['asc', 'desc'] suggestions
  * - `percentile`: Marks 'percentile' parameter as constantOnly
+ * - `count_distinct`: Marks 'precision' parameter as constantOnly
  * - `count`: Marks 'percentile' parameter as constantOnly
+ * - `round`: Marks 'decimals' parameter as constantOnly
+ * - `round_to`: Marks 'points' parameter as constantOnly
  * - `qstr`: Adds custom parameter snippet for triple-quoted strings
  * - `kql`: Adds custom parameter snippet for triple-quoted strings
  */
@@ -161,6 +164,12 @@ export function enrichFunctionParameters(functionDefinition: FunctionDefinition)
     });
   }
 
+  if (functionDefinition.name === 'count_distinct') {
+    return enrichFunctionSignatures(functionDefinition, 'precision', {
+      constantOnly: true,
+    });
+  }
+
   if (functionDefinition.name === 'count') {
     return enrichFunctionSignatures(functionDefinition, 'percentile', {
       constantOnly: true,
@@ -173,6 +182,18 @@ export function enrichFunctionParameters(functionDefinition: FunctionDefinition)
     });
   }
 
+  if (functionDefinition.name === 'round') {
+    return enrichFunctionSignatures(functionDefinition, 'decimals', {
+      constantOnly: true,
+    });
+  }
+
+  if (functionDefinition.name === 'round_to') {
+    return enrichFunctionSignatures(functionDefinition, 'points', {
+      constantOnly: true,
+    });
+  }
+
   if (functionDefinition.name === 'qstr') {
     return {
       ...functionDefinition,

src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts

@@ -24,7 +24,15 @@ import { aggFunctionDefinitions } from '../definitions/generated/aggregation_fun
 import { timeSeriesAggFunctionDefinitions } from '../definitions/generated/time_series_agg_functions';
 import { groupingFunctionDefinitions } from '../definitions/generated/grouping_functions';
 import { scalarFunctionDefinitions } from '../definitions/generated/scalar_functions';
-import { operatorsDefinitions } from '../definitions/all_operators';
+import {
+  operatorsDefinitions,
+  comparisonFunctions,
+  logicalOperators,
+  arithmeticOperators,
+  patternMatchOperators,
+  inOperators,
+  nullCheckOperators,
+} from '../definitions/all_operators';
 import { parse } from '../parser';
 import type { ESQLAstAllCommands } from '../types';
 import type {
@@ -164,6 +172,7 @@ export function getFunctionSignaturesByReturnType(
     grouping,
     scalar,
     operators,
+    excludeOperatorGroups,
     // skipAssign here is used to communicate to not propose an assignment if it's not possible
     // within the current context (the actual logic has it, but here we want a shortcut)
     skipAssign,
@@ -173,6 +182,10 @@ export function getFunctionSignaturesByReturnType(
     grouping?: boolean;
     scalar?: boolean;
     operators?: boolean;
+    /** Exclude specific operator groups (e.g., ['in', 'nullCheck']) */
+    excludeOperatorGroups?: Array<
+      'logical' | 'comparison' | 'arithmetic' | 'pattern' | 'in' | 'nullCheck'
+    >;
     skipAssign?: boolean;
   } = {},
   paramsTypes?: Readonly<FunctionParameterType[]>,
@@ -201,7 +214,50 @@ export function getFunctionSignaturesByReturnType(
     list.push(...scalarFunctionDefinitions);
   }
   if (operators) {
-    list.push(...operatorsDefinitions.filter(({ name }) => (skipAssign ? name !== '=' : true)));
+    const hasStringParams = paramsTypes?.some((type) => type === 'text' || type === 'keyword');
+    const comparisonOperatorNames = comparisonFunctions.map(({ name }) => name);
+
+    // Build set of operator names to exclude based on excludeOperatorGroups
+    const excludedOperatorNames = new Set<string>();
+
+    if (excludeOperatorGroups) {
+      const operatorGroupMap: Record<string, Array<{ name: string; signatures?: unknown[] }>> = {
+        logical: logicalOperators,
+        comparison: comparisonFunctions,
+        arithmetic: arithmeticOperators,
+        pattern: patternMatchOperators,
+        in: inOperators,
+        nullCheck: nullCheckOperators,
+      };
+
+      for (const groupName of excludeOperatorGroups) {
+        const group = operatorGroupMap[groupName];
+
+        if (group) {
+          for (const op of group) {
+            excludedOperatorNames.add(op.name);
+          }
+        }
+      }
+    }
+
+    const filteredOperators = operatorsDefinitions.filter(({ name }) => {
+      if (skipAssign && (name === '=' || name === ':')) {
+        return false;
+      }
+
+      if (hasStringParams && comparisonOperatorNames.includes(name)) {
+        return false;
+      }
+
+      if (excludedOperatorNames.has(name)) {
+        return false;
+      }
+
+      return true;
+    });
+
+    list.push(...filteredOperators);
   }
 
   const deduped = Array.from(new Set(list));
@@ -299,3 +355,17 @@ export const containsSnippet = (text: string): boolean => {
   const snippetRegex = /\$(\d+|\{\d+:[^}]*\})/;
   return snippetRegex.test(text);
 };
+
+/**
+ * Convert operator definition groups to suggestion strings.
+ * Use this instead of hardcoding operator strings in tests.
+ */
+export function getOperatorSuggestions(
+  operators: Array<{ name: string; signatures: Array<{ params: unknown[] }> }>
+): string[] {
+  return operators.map(({ name, signatures }) =>
+    signatures.some(({ params }) => params.length > 1)
+      ? `${name.toUpperCase()} $0`
+      : name.toUpperCase()
+  );
+}

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.test.ts

@@ -13,10 +13,19 @@ import {
   expectSuggestions,
   getFieldNamesByType,
   getFunctionSignaturesByReturnType,
+  getOperatorSuggestions,
 } from '../../../__tests__/autocomplete';
 import type { ICommandCallbacks } from '../../types';
 import { ESQL_COMMON_NUMERIC_TYPES } from '../../../definitions/types';
 import { timeUnitsToSuggest } from '../../../definitions/constants';
+import {
+  logicalOperators,
+  arithmeticOperators,
+  comparisonFunctions,
+  patternMatchOperators,
+  inOperators,
+  nullCheckOperators,
+} from '../../../definitions/all_operators';
 
 const roundParameterTypes = ['double', 'integer', 'long', 'unsigned_long'] as const;
 
@@ -177,7 +186,9 @@ describe('EVAL Autocomplete', () => {
     );
     await evalExpectSuggestions(
       'from index | EVAL keywordField not ',
-      ['LIKE $0', 'RLIKE $0', 'IN $0'],
+      getOperatorSuggestions(
+        [...inOperators, ...patternMatchOperators].filter((op) => !op.name.startsWith('not '))
+      ),
       mockCallbacks
     );
 
@@ -223,10 +234,7 @@ describe('EVAL Autocomplete', () => {
           { operators: true, skipAssign: true },
           ['double', 'long']
         ),
-        'IN $0',
-        'IS NOT NULL',
-        'IS NULL',
-        'NOT IN $0',
+        ...getOperatorSuggestions([...inOperators, ...nullCheckOperators]),
       ],
       mockCallbacks
     );
@@ -250,34 +258,8 @@ describe('EVAL Autocomplete', () => {
     (mockCallbacks.getByType as jest.Mock).mockResolvedValue(
       expectedDoubleIntegerFields.map((name) => ({ label: name, text: name }))
     );
-    await evalExpectSuggestions(
-      'from a | eval a=round(doubleField, ',
-      [
-        ...expectedDoubleIntegerFields,
-        ...getFunctionSignaturesByReturnType(
-          Location.EVAL,
-          ['integer', 'long'],
-          { scalar: true },
-          undefined,
-          ['round']
-        ),
-      ],
-      mockCallbacks
-    );
-    await evalExpectSuggestions(
-      'from a | eval round(doubleField, ',
-      [
-        ...expectedDoubleIntegerFields,
-        ...getFunctionSignaturesByReturnType(
-          Location.EVAL,
-          ['integer', 'long'],
-          { scalar: true },
-          undefined,
-          ['round']
-        ),
-      ],
-      mockCallbacks
-    );
+    await evalExpectSuggestions('from a | eval a=round(doubleField, ', [], mockCallbacks);
+    await evalExpectSuggestions('from a | eval round(doubleField, ', [], mockCallbacks);
     const expectedAny = getFieldNamesByType('any');
     (mockCallbacks.getByType as jest.Mock).mockResolvedValue(
       expectedAny.map((name) => ({ label: name, text: name }))
@@ -506,32 +488,34 @@ describe('EVAL Autocomplete', () => {
 
     test('suggests operators after initial column based on type', async () => {
       // case( field ) suggests all appropriate operators for that field type
+      // Note: CASE is expression-heavy, comma is not automatically suggested after fields
       await evalExpectSuggestions('from a | eval case( textField ', [
-        ...getFunctionSignaturesByReturnType(
-          Location.EVAL,
-          'any',
-          { operators: true, skipAssign: true, agg: false, scalar: false },
-          ['text']
-        ),
+        ...getOperatorSuggestions([
+          ...patternMatchOperators,
+          ...inOperators,
+          ...nullCheckOperators,
+        ]),
       ]);
 
       await evalExpectSuggestions('from a | eval case( doubleField ', [
-        ...getFunctionSignaturesByReturnType(
-          Location.EVAL,
-          'any',
-          { operators: true, skipAssign: true, agg: false, scalar: false },
-          ['double']
-        ),
+        ...getOperatorSuggestions([
+          ...comparisonFunctions,
+          // Exclude unary-only operators (negation) - autocomplete doesn't suggest them
+          ...arithmeticOperators.filter(
+            (op) => !op.signatures.every((sig) => sig.params.length === 1)
+          ),
+          ...inOperators,
+          ...nullCheckOperators,
+        ]),
       ]);
 
       await evalExpectSuggestions('from a | eval case( booleanField ', [
-        ...getFunctionSignaturesByReturnType(
-          Location.EVAL,
-          'any',
-          { operators: true, skipAssign: true, agg: false, scalar: false },
-          ['boolean']
+        ...getOperatorSuggestions(
+          comparisonFunctions.filter(({ name }) => name === '==' || name === '!=')
         ),
         ',',
+        ...getOperatorSuggestions([...logicalOperators, ...inOperators, ...nullCheckOperators]),
+        'NOT',
       ]);
     });
 

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts

@@ -62,6 +62,9 @@ export async function autocomplete(
     location: Location.EVAL,
     context,
     callbacks,
+    options: {
+      preferredExpressionType: 'any',
+    },
   });
 
   const positionInExpression = getExpressionPosition(query, expressionRoot);

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/autocomplete.test.ts

@@ -369,7 +369,7 @@ describe('FORK Autocomplete', () => {
                 ...getFunctionSignaturesByReturnType(
                   Location.STATS_WHERE,
                   'any',
-                  { operators: true },
+                  { operators: true, skipAssign: true },
                   ['integer']
                 ),
               ]

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/autocomplete.test.ts

@@ -22,17 +22,24 @@ import {
 import { expectSuggestions, suggest } from '../../../__tests__/autocomplete';
 import type { ICommandCallbacks } from '../../types';
 import { buildConstantsDefinitions } from '../../../definitions/utils/literals';
+import {
+  logicalOperators,
+  comparisonFunctions,
+  patternMatchOperators,
+  inOperators,
+  nullCheckOperators,
+} from '../../../definitions/all_operators';
 
 // ============================================================================
-// Single Source of Truth - Operator Suggestions by Group
+// Operator Suggestions - Derived from Real Definitions
 // ============================================================================
 
 const OPERATOR_SUGGESTIONS = {
-  LOGICAL: ['AND', 'OR'],
-  COMPARISON: ['!=', '==', '>', '>=', '<', '<='],
-  PATTERN: ['LIKE', 'NOT LIKE', 'RLIKE', 'NOT RLIKE', ':'],
-  SET: ['IN', 'NOT IN'],
-  EXISTENCE: ['IS NULL', 'IS NOT NULL'],
+  LOGICAL: logicalOperators.map(({ name }) => name.toUpperCase()),
+  COMPARISON: comparisonFunctions.map(({ name }) => name.toUpperCase()),
+  PATTERN: [...patternMatchOperators.map(({ name }) => name.toUpperCase()), ':'], // ':' is assignment, keep for now
+  SET: inOperators.map(({ name }) => name.toUpperCase()),
+  EXISTENCE: nullCheckOperators.map(({ name }) => name.toUpperCase()),
 };
 
 // Helper to add placeholder to operator labels for test expectations
@@ -339,9 +346,13 @@ describe('RERANK Autocomplete', () => {
       await expectRerankSuggestions(query, {
         contains: [
           addPlaceholder([OPERATOR_SUGGESTIONS.PATTERN[0]])[0],
-          ...addPlaceholder(OPERATOR_SUGGESTIONS.COMPARISON.slice(0, 2)),
+          ...addPlaceholder(OPERATOR_SUGGESTIONS.SET.slice(0, 1)),
+          OPERATOR_SUGGESTIONS.EXISTENCE[0],
+        ],
+        notContains: [
+          ...NEXT_ACTIONS_EXPRESSIONS,
+          ...addPlaceholder(OPERATOR_SUGGESTIONS.COMPARISON),
         ],
-        notContains: NEXT_ACTIONS_EXPRESSIONS,
       });
     });
   });