[ES|QL] Rerank command validation (#233691)

## Summary

Part of #217285

We enable validation for the RERANK command

Many validations come from parsing errors and are therefore difficult to
read.
- The fragment after ON should be standard for other commands and
therefore no custom handlers.
- WITH needs to be handled custom, like COMPLETION.
- I have a custom handling for the query string value similar to
COMPLETION.

Note: For RERANK, our validation only works after the ON command is
recognized, because the grammar requires it to define the input. Until
the parser consumes On, the builder has no scope and we can see only
syntax errors.

<img width="1304" height="123" alt="Screenshot 2025-09-02 080759"
src="https://github.com/user-attachments/assets/eb03e855-e957-41a0-8dff-4cec6e1ed9f0"
/>
<img width="916" height="332" alt="Screenshot 2025-09-02 080925"
src="https://github.com/user-attachments/assets/859b97f9-054e-4c26-99f3-45ec8c3d8672"
/>
<img width="719" height="125" alt="Screenshot 2025-09-02 081011"
src="https://github.com/user-attachments/assets/28f3ecc3-6d64-4850-b412-2939553faf3a"
/>

<img width="1301" height="122" alt="Screenshot 2025-09-02 081044"
src="https://github.com/user-attachments/assets/145cb625-e9ad-418d-a0bf-fbb0d98b8b1d"
/>

Author: bartoval

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.test.ts

@@ -104,19 +104,19 @@ describe('COMPLETION Validation', () => {
       it('prompt is a constant, but not text', async () => {
         await completionExpectErrors(
           `FROM index | COMPLETION 47 WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [integer]']
+          ['COMPLETION query must be of type text. Found integer']
         );
 
         await completionExpectErrors(
           `FROM index | COMPLETION true WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [boolean]']
+          ['COMPLETION query must be of type text. Found boolean']
         );
       });
 
       it('prompt is a function, but does not return text', async () => {
         await completionExpectErrors(
           `FROM index | COMPLETION PI() WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [double]']
+          ['COMPLETION query must be of type text. Found double']
         );
 
         await completionExpectErrors(
@@ -126,18 +126,18 @@ describe('COMPLETION Validation', () => {
                       var0 == 0, 2,
                       var0 == 1, 1
                     ) WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [integer]']
+          ['COMPLETION query must be of type text. Found integer']
         );
 
         await completionExpectErrors(
           `FROM index | COMPLETION TO_DATETIME("2023-12-02T11:00:00.000Z") WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [date]']
+          ['COMPLETION query must be of type text. Found date']
         );
 
         await completionExpectErrors(
           `FROM index | COMPLETION AVG(integerField) WITH { "inference_id": "inferenceId"}`,
           [
-            '[COMPLETION] prompt must be of type [text] but is [double]',
+            'COMPLETION query must be of type text. Found double',
             'Function AVG not allowed in COMPLETION',
           ]
         );
@@ -146,19 +146,19 @@ describe('COMPLETION Validation', () => {
       it('prompt is a column, but not a text one', () => {
         completionExpectErrors(
           `FROM index | EVAL integerPrompt = 47 | COMPLETION integerPrompt WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [integer]']
+          ['COMPLETION query must be of type text. Found integer']
         );
         completionExpectErrors(
           `FROM index | EVAL ipPrompt = to_ip("1.2.3.4") | COMPLETION ipPrompt WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [ip]']
+          ['COMPLETION query must be of type text. Found ip']
         );
         completionExpectErrors(
           `FROM index | COMPLETION dateField WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [date]']
+          ['COMPLETION query must be of type text. Found date']
         );
         completionExpectErrors(
           `FROM index | COMPLETION counterIntegerField WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [counter_integer]']
+          ['COMPLETION query must be of type text. Found counter_integer']
         );
       });
 
@@ -176,26 +176,26 @@ describe('COMPLETION Validation', () => {
       it('prompt is a parenthesized expression, but not a text one', () => {
         completionExpectErrors(
           `FROM index | COMPLETION (1 > 2) WITH { "inference_id": "inferenceId"}`,
-          ['[COMPLETION] prompt must be of type [text] but is [boolean]']
+          ['COMPLETION query must be of type text. Found boolean']
         );
       });
 
       it('inference_id is not provided', () => {
         completionExpectErrors(`FROM index | COMPLETION "prompt"`, [
-          '[COMPLETION] inference_id parameter is required',
+          '"inference_id" parameter is required for COMPLETION.',
         ]);
         completionExpectErrors(`FROM index | COMPLETION "prompt" WITH`, [
-          '[COMPLETION] inference_id parameter is required',
+          '"inference_id" parameter is required for COMPLETION.',
         ]);
         completionExpectErrors(`FROM index | COMPLETION "prompt" WITH {}`, [
-          '[COMPLETION] inference_id parameter is required',
+          '"inference_id" parameter is required for COMPLETION.',
         ]);
         completionExpectErrors(`FROM index | COMPLETION "prompt" WITH { "": ""}`, [
-          '[COMPLETION] inference_id parameter is required',
+          '"inference_id" parameter is required for COMPLETION.',
         ]);
         completionExpectErrors(
           `FROM index | COMPLETION "prompt" WITH { "some_param": "some_value"}`,
-          ['[COMPLETION] inference_id parameter is required']
+          ['"inference_id" parameter is required for COMPLETION.']
         );
       });
     });

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts

@@ -6,11 +6,11 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { i18n } from '@kbn/i18n';
 import type { ESQLAst, ESQLAstCompletionCommand, ESQLCommand, ESQLMessage } from '../../../types';
 import type { ICommandContext, ICommandCallbacks } from '../../types';
 import { getExpressionType } from '../../../definitions/utils/expressions';
 import { validateCommandArguments } from '../../../definitions/utils/validation';
+import { errors } from '../../../definitions/utils/errors';
 
 const supportedPromptTypes = ['text', 'keyword', 'unknown', 'param'];
 
@@ -31,27 +31,16 @@ export const validate = (
   );
 
   if (!supportedPromptTypes.includes(promptExpressionType)) {
-    messages.push({
-      location: 'location' in prompt ? prompt?.location : location,
-      text: i18n.translate('kbn-esql-ast.esql.validation.completionUnsupportedFieldType', {
-        defaultMessage:
-          '[COMPLETION] prompt must be of type [text] but is [{promptExpressionType}]',
-        values: { promptExpressionType },
-      }),
-      type: 'error',
-      code: 'completionUnsupportedFieldType',
-    });
+    messages.push(
+      errors.byId('unsupportedQueryType', 'location' in prompt ? prompt?.location : location, {
+        command: 'COMPLETION',
+        expressionType: promptExpressionType,
+      })
+    );
   }
 
-  if (inferenceId.incomplete) {
-    messages.push({
-      location: command.location,
-      text: i18n.translate('kbn-esql-ast.esql.validation.completionInferenceIdRequired', {
-        defaultMessage: '[COMPLETION] inference_id parameter is required',
-      }),
-      type: 'error',
-      code: 'completionInferenceIdRequired',
-    });
+  if (inferenceId?.incomplete) {
+    messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'COMPLETION' }));
   }
 
   const targetName = targetField?.name || 'completion';

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { validate } from './validate';
+import { expectErrors } from '../../../__tests__/validation';
+import { mockContext } from '../../../__tests__/context_fixtures';
+
+const rerankExpectErrors = (query: string, expected: string[]) =>
+  expectErrors(query, expected, mockContext, 'rerank', validate);
+
+describe('RERANK Validation', () => {
+  describe('Query', () => {
+    test('query as parameter', () => {
+      rerankExpectErrors('FROM index | RERANK ?q ON keywordField', []);
+    });
+
+    test('unsupported field type', () => {
+      rerankExpectErrors('FROM index | RERANK col0=2 ON keywordField', [
+        'RERANK query must be of type text. Found integer',
+      ]);
+    });
+  });
+
+  describe('Inference_id check', () => {
+    test('inference_id as parameter', () => {
+      rerankExpectErrors(
+        'FROM index | RERANK "q" ON keywordField WITH { "inference_id": ?named_param }',
+        []
+      );
+    });
+
+    const msg = '"inference_id" parameter is required for RERANK.';
+
+    test('WITH without map (partial: WITH)', () => {
+      rerankExpectErrors('FROM index | RERANK "q" ON keywordField WITH', [msg]);
+    });
+
+    test('WITH with partial map start (WITH{)', () => {
+      rerankExpectErrors('FROM index | RERANK "q" ON keywordField WITH{', [msg]);
+    });
+
+    test('WITH with empty map (WITH {})', () => {
+      rerankExpectErrors('FROM index | RERANK "q" ON keywordField WITH { }', [msg]);
+    });
+
+    test('WITH map without inference_id key', () => {
+      rerankExpectErrors('FROM index | RERANK "q" ON keywordField WITH { "some_param": "value" }', [
+        msg,
+      ]);
+    });
+
+    test('WITH inference_id with empty string value', () => {
+      rerankExpectErrors('FROM index | RERANK "q" ON keywordField WITH { "inference_id": "" }', [
+        msg,
+      ]);
+    });
+  });
+});

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

@@ -7,10 +7,13 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { i18n } from '@kbn/i18n';
 import type { ESQLCommand, ESQLMessage, ESQLAst, ESQLAstRerankCommand } from '../../../types';
 import type { ICommandContext, ICommandCallbacks } from '../../types';
+import { getExpressionType } from '../../../definitions/utils/expressions';
 import { validateCommandArguments } from '../../../definitions/utils/validation';
+import { errors } from '../../../definitions/utils/errors';
+
+const supportedQueryTypes = ['keyword', 'text', 'param'];
 
 export const validate = (
   command: ESQLCommand,
@@ -20,35 +23,39 @@ export const validate = (
 ): ESQLMessage[] => {
   const messages: ESQLMessage[] = [];
 
-  // Run standard argument validation
-  messages.push(...validateCommandArguments(command, ast, context, callbacks));
+  const { query, targetField, location, inferenceId } = command as ESQLAstRerankCommand;
+  const rerankExpressionType = getExpressionType(
+    query,
+    context?.fields,
+    context?.userDefinedColumns
+  );
 
-  // Cast to RERANK command for type-specific validation
-  const rerankCommand = command as ESQLAstRerankCommand;
-
-  // Validate that query text is provided
-  if (!rerankCommand.query) {
-    messages.push({
-      location: command.location,
-      text: i18n.translate('kbn-esql-ast.esql.validation.rerankMissingQuery', {
-        defaultMessage: '[RERANK] Query text is required.',
-      }),
-      type: 'error',
-      code: 'rerankMissingQuery',
-    });
+  // check for supported query types
+  if (!supportedQueryTypes.includes(rerankExpressionType)) {
+    messages.push(
+      errors.byId('unsupportedQueryType', 'location' in query ? query?.location : location, {
+        command: 'RERANK',
+        expressionType: rerankExpressionType,
+      })
+    );
   }
 
-  // Validate that at least one field is specified in ON clause
-  if (!rerankCommand.fields || rerankCommand.fields.length === 0) {
-    messages.push({
-      location: command.location,
-      text: i18n.translate('kbn-esql-ast.esql.validation.rerankMissingFields', {
-        defaultMessage: '[RERANK] At least one field must be specified in the ON clause.',
-      }),
-      type: 'error',
-      code: 'rerankMissingFields',
-    });
+  if (inferenceId?.incomplete) {
+    messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'RERANK' }));
   }
 
+  const targetName = targetField?.name || 'rerank';
+
+  // Sets the target field so the column is recognized after the command is applied
+  context?.userDefinedColumns.set(targetName, [
+    {
+      name: targetName,
+      location: targetField?.location || location,
+      type: 'keyword',
+    },
+  ]);
+
+  messages.push(...validateCommandArguments(command, ast, context, callbacks));
+
   return messages;
 };

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts

@@ -213,6 +213,11 @@ export enum Location {
    */
   RENAME = 'rename',
 
+  /**
+   * In the RERANK command
+   */
+  RERANK = 'rerank',
+
   /**
    * In the JOIN command (used only for AS)
    */
@@ -244,6 +249,7 @@ const commandOptionNameToLocation: Record<string, Location> = {
   join: Location.JOIN,
   show: Location.SHOW,
   completion: Location.COMPLETION,
+  rerank: Location.RERANK,
 };
 
 /**

src/platform/packages/shared/kbn-esql-ast/src/definitions/all_operators.ts

@@ -95,6 +95,7 @@ const otherDefinitions: FunctionDefinition[] = [
       Location.DISSECT,
       Location.COMPLETION,
       Location.RENAME,
+      Location.RERANK,
     ],
     signatures: [
       {

src/platform/packages/shared/kbn-esql-ast/src/definitions/types.ts

@@ -356,6 +356,14 @@ export interface ValidationErrors {
     message: string;
     type: {};
   };
+  inferenceIdRequired: {
+    message: string;
+    type: {};
+  };
+  unsupportedQueryType: {
+    message: string;
+    type: {};
+  };
 }
 
 export type ErrorTypes = keyof ValidationErrors;

src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/errors.ts

@@ -294,6 +294,24 @@ Expected one of:
         }),
         type: 'warning',
       };
+
+    case 'inferenceIdRequired':
+      return {
+        message: i18n.translate('kbn-esql-ast.esql.validation.inferenceIdRequired', {
+          defaultMessage: '"inference_id" parameter is required for {command}.',
+          values: { command: out.command.toUpperCase() },
+        }),
+        type: 'error',
+      };
+
+    case 'unsupportedQueryType':
+      return {
+        message: i18n.translate('kbn-esql-ast.esql.validation.unsupportedQueryType', {
+          defaultMessage: '{command} query must be of type text. Found {expressionType}',
+          values: { command: out.command.toUpperCase(), expressionType: out.expressionType },
+        }),
+        type: 'error',
+      };
   }
   return { message: '' };
 }

src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/rerank.test.ts

@@ -200,7 +200,7 @@ describe('RERANK', () => {
         fields: [{ type: 'column', name: 'title' }],
       });
 
-      expect(rerankCmd).not.toHaveProperty('inferenceId');
+      expect(rerankCmd.inferenceId).toEqual(undefined);
     });
 
     it('should handle missing ON clause', () => {