[ES|QL] Add support for RRF

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

@@ -0,0 +1,31 @@
+/*
+ * 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 { parse } from '../parser';
+
+describe('RRF', () => {
+  describe('correctly formatted', () => {
+    it('can parse RRF command without modifiers', () => {
+      const text = `FROM search-movies METADATA _score, _id, _index
+| FORK ( WHERE semantic_title:"Shakespeare" | SORT _score)
+          ( WHERE title:"Shakespeare" | SORT _score)
+| RRF
+| KEEP title, _score`;
+
+      const { root, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(root.commands[2]).toMatchObject({
+        type: 'command',
+        name: 'rrf',
+        args: [],
+      });
+    });
+  });
+});

src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts

@@ -32,6 +32,7 @@ import {
   type TimeSeriesCommandContext,
   type WhereCommandContext,
   RerankCommandContext,
+  RrfCommandContext,
 } from '../antlr/esql_parser';
 import { default as ESQLParserListener } from '../antlr/esql_parser_listener';
 import type { ESQLAst } from '../types';
@@ -351,6 +352,20 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
     this.ast.push(command);
   }
 
+  /**
+   * Exit a parse tree produced by `esql_parser.rrfCommand`.
+   *
+   * Parse the RRF (Reciprocal Rank Fusion) command:
+   *
+   * RRF
+   *
+   * @param ctx the parse tree
+   */
+  exitRrfCommand(ctx: RrfCommandContext): void {
+    const command = createCommand('rrf', ctx);
+    this.ast.push(command);
+  }
+
   enterEveryRule(ctx: ParserRuleContext): void {
     // method not implemented, added to satisfy interface expectation
   }

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts

@@ -235,6 +235,32 @@ describe('single line query', () => {
         );
       });
     });
+
+    describe('RRF', () => {
+      test('from single line', () => {
+        const { text } =
+          reprint(`FROM search-movies METADATA _score, _id, _index | FORK (WHERE semantic_title : "Shakespeare" | SORT _score) (WHERE title : "Shakespeare" | SORT _score) | RRF | KEEP title, _score
+        `);
+
+        expect(text).toBe(
+          'FROM search-movies METADATA _score, _id, _index | FORK (WHERE semantic_title : "Shakespeare" | SORT _score) (WHERE title : "Shakespeare" | SORT _score) | RRF | KEEP title, _score'
+        );
+      });
+
+      test('from multiline', () => {
+        const { text } = reprint(`FROM search-movies METADATA _score, _id, _index
+  | FORK
+    (WHERE semantic_title : "Shakespeare" | SORT _score)
+    (WHERE title : "Shakespeare" | SORT _score)
+  | RRF
+  | KEEP title, _score
+          `);
+
+        expect(text).toBe(
+          'FROM search-movies METADATA _score, _id, _index | FORK (WHERE semantic_title : "Shakespeare" | SORT _score) (WHERE title : "Shakespeare" | SORT _score) | RRF | KEEP title, _score'
+        );
+      });
+    });
   });
 
   describe('expressions', () => {

src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts

@@ -520,6 +520,12 @@ export class ForkCommandVisitorContext<
   Data extends SharedData = SharedData
 > extends CommandVisitorContext<Methods, Data, ESQLAstCommand> {}
 
+// RRF
+export class RrfCommandVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData
+> extends CommandVisitorContext<Methods, Data, ESQLAstCommand> {}
+
 // Expressions -----------------------------------------------------------------
 
 export class ExpressionVisitorContext<

src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts

@@ -197,6 +197,10 @@ export class GlobalVisitorContext<
         if (!this.methods.visitForkCommand) break;
         return this.visitForkCommand(parent, commandNode, input as any);
       }
+      case 'rrf': {
+        if (!this.methods.visitRrfCommand) break;
+        return this.visitRrfCommand(parent, commandNode, input as any);
+      }
     }
     return this.visitCommandGeneric(parent, commandNode, input as any);
   }
@@ -417,6 +421,15 @@ export class GlobalVisitorContext<
     return this.visitWithSpecificContext('visitForkCommand', context, input);
   }
 
+  public visitRrfCommand(
+    parent: contexts.VisitorContext | null,
+    node: ESQLAstCommand,
+    input: types.VisitorInput<Methods, 'visitRrfCommand'>
+  ): types.VisitorOutput<Methods, 'visitRrfCommand'> {
+    const context = new contexts.RrfCommandVisitorContext(this, node, parent);
+    return this.visitWithSpecificContext('visitRrfCommand', context, input);
+  }
+
   // #endregion
 
   // #region Expression visiting -------------------------------------------------------

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

@@ -189,6 +189,7 @@ export interface VisitorMethods<
   >;
   visitForkCommand?: Visitor<contexts.ForkCommandVisitorContext<Visitors, Data>, any, any>;
   visitCommandOption?: Visitor<contexts.CommandOptionVisitorContext<Visitors, Data>, any, any>;
+  visitRrfCommand?: Visitor<contexts.RrfCommandVisitorContext<Visitors, Data>, any, any>;
   visitExpression?: Visitor<contexts.ExpressionVisitorContext<Visitors, Data>, any, any>;
   visitSourceExpression?: Visitor<
     contexts.SourceExpressionVisitorContext<Visitors, Data>,

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

@@ -13,6 +13,14 @@ import { ESQLFieldWithMetadata } from '../validation/types';
 import { fieldTypes } from '../definitions/types';
 import { ESQLCallbacks } from '../shared/types';
 
+export const metadataFields: ESQLFieldWithMetadata[] = [
+  { name: '_id', type: 'keyword' },
+  { name: '_index', type: 'keyword' },
+  { name: '_score', type: 'double' },
+  { name: '_type', type: 'keyword' },
+  { name: '_source', type: 'keyword' },
+];
+
 export const fields: ESQLFieldWithMetadata[] = [
   ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type })),
   { name: 'any#Char$Field', type: 'double' },
@@ -107,6 +115,9 @@ export function getCallbackMocks(): ESQLCallbacks {
         };
         return [field];
       }
+      if (/METADATA/.test(query)) {
+        return [...fields, ...metadataFields];
+      }
       return fields;
     }),
     getSources: jest.fn(async () =>

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.rff.test.ts

@@ -0,0 +1,41 @@
+/*
+ * 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 { setup } from './helpers';
+
+describe('autocomplete.suggest', () => {
+  describe('RRF', () => {
+    it('does not suggest RRF in the general list of commands', async () => {
+      const { suggest } = await setup();
+      const suggestedCommands = (await suggest('FROM index | /')).map((s) => s.text);
+      expect(suggestedCommands).not.toContain('RRF ');
+    });
+
+    it('suggests RRF immediately after a FORK command', async () => {
+      const { suggest } = await setup();
+      const suggestedCommands = (await suggest('FROM a | FORK (LIMIT 1) (LIMIT 2) | /')).map(
+        (s) => s.text
+      );
+      expect(suggestedCommands).toContain('RRF ');
+    });
+
+    it('does not suggests RRF if FORK is not immediately before', async () => {
+      const { suggest } = await setup();
+      const suggestedCommands = (
+        await suggest('FROM a | FORK (LIMIT 1) (LIMIT 2) | LIMIT 1 /')
+      ).map((s) => s.text);
+      expect(suggestedCommands).not.toContain('RRF ');
+    });
+
+    it('suggests pipe after complete command', async () => {
+      const { assertSuggestions } = await setup();
+      await assertSuggestions('FROM a | FORK (LIMIT 1) (LIMIT 2) | RRF /', ['| ']);
+    });
+  });
+});

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts

@@ -118,8 +118,9 @@ export async function suggest(
 
   if (astContext.type === 'newCommand') {
     // propose main commands here
+    // resolve particular commands suggestions after
     // filter source commands if already defined
-    const suggestions = getCommandAutocompleteDefinitions(getAllCommands());
+    let suggestions = getCommandAutocompleteDefinitions(getAllCommands());
     if (!ast.length) {
       // Display the recommended queries if there are no commands (empty state)
       const recommendedQueriesSuggestions: SuggestionRawDefinition[] = [];
@@ -144,6 +145,13 @@ export async function suggest(
       return [...sourceCommandsSuggestions, ...recommendedQueriesSuggestions];
     }
 
+    const lastCommand = root.commands[root.commands.length - 1];
+    const lastCommandDefinition = getCommandDefinition(lastCommand.name);
+
+    if (lastCommandDefinition.commandsSuggestionsAfter) {
+      suggestions = lastCommandDefinition.commandsSuggestionsAfter(suggestions);
+    }
+
     return suggestions.filter((def) => !isSourceCommand(def));
   }
 

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/fork/commands_suggestions_after.ts

@@ -0,0 +1,24 @@
+/*
+ * 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 { CommandDefinition, getCommandDefinition } from '../../../..';
+import { getCommandAutocompleteDefinitions } from '../../complete_items';
+import { SuggestionRawDefinition } from '../../types';
+
+export function commandsSuggestionsAfter(suggestions: SuggestionRawDefinition[]) {
+  return [
+    ...suggestions,
+    ...getCommandAutocompleteDefinitions([
+      {
+        ...(getCommandDefinition('rrf') as CommandDefinition<string>),
+        hidden: false,
+      },
+    ]),
+  ];
+}

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/fork/index.ts

@@ -9,3 +9,4 @@
 
 export { suggest } from './suggest';
 export { fieldsSuggestionsAfter } from './fields_suggestions_after';
+export { commandsSuggestionsAfter } from './commands_suggestions_after';

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/rrf/index.ts

@@ -0,0 +1,16 @@
+/*
+ * 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 { CommandSuggestParams } from '../../../definitions/types';
+import type { SuggestionRawDefinition } from '../../types';
+import { pipeCompleteItem } from '../../complete_items';
+
+export function suggest(_params: CommandSuggestParams<'rrf'>): SuggestionRawDefinition[] {
+  return [pipeCompleteItem];
+}

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/definitions/commands.ts

@@ -51,6 +51,7 @@ import { suggest as suggestForEval } from '../autocomplete/commands/eval';
 import {
   suggest as suggestForFork,
   fieldsSuggestionsAfter as fieldsSuggestionsAfterFork,
+  commandsSuggestionsAfter as commandsSuggestionsAfterFork,
 } from '../autocomplete/commands/fork';
 import { suggest as suggestForFrom } from '../autocomplete/commands/from';
 import { suggest as suggestForTimeseries } from '../autocomplete/commands/timeseries';
@@ -69,6 +70,8 @@ import {
   suggest as suggestForRename,
   fieldsSuggestionsAfter as fieldsSuggestionsAfterRename,
 } from '../autocomplete/commands/rename';
+import { suggest as suggestForRrf } from '../autocomplete/commands/rrf';
+import { validate as validateRrf } from '../validation/commands/rrf';
 import { suggest as suggestForRow } from '../autocomplete/commands/row';
 import { suggest as suggestForShow } from '../autocomplete/commands/show';
 import { suggest as suggestForSort } from '../autocomplete/commands/sort';
@@ -705,5 +708,19 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
     },
 
     fieldsSuggestionsAfter: fieldsSuggestionsAfterFork,
+    commandsSuggestionsAfter: commandsSuggestionsAfterFork,
+  },
+  {
+    hidden: true, // This command should be keep as hidden as it only must be shown after a FORK command
+    preview: true,
+    name: 'rrf',
+    description: i18n.translate('kbn-esql-validation-autocomplete.esql.definitions.rrfDoc', {
+      defaultMessage:
+        'Combines multiple result sets with different scoring functions into a single result set.',
+    }),
+    declaration: `RRF`, // TODO: update when options are added
+    examples: ['… FORK (LIMIT 1) (LIMIT 2) | RRF'],
+    suggest: suggestForRrf,
+    validate: validateRrf,
   },
 ];

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

@@ -6,13 +6,14 @@
  * 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 type {
+import {
   ESQLAstItem,
   ESQLCommand,
   ESQLFunction,
   ESQLMessage,
   ESQLSource,
   ESQLAstCommand,
+  ESQLAst,
 } from '@kbn/esql-ast';
 import { ESQLControlVariable } from '@kbn/esql-types';
 import { GetColumnsByTypeFn, SuggestionRawDefinition } from '../autocomplete/types';
@@ -419,7 +420,11 @@ export interface CommandDefinition<CommandName extends string> {
    * prevent the default behavior. If you need a full override, we are currently
    * doing those directly in the validateCommand function in the validation module.
    */
-  validate?: (command: ESQLCommand<CommandName>, references: ReferenceMaps) => ESQLMessage[];
+  validate?: (
+    command: ESQLCommand<CommandName>,
+    references: ReferenceMaps,
+    ast: ESQLAst
+  ) => ESQLMessage[];
 
   /**
    * This method is called to load suggestions when the cursor is within this command.
@@ -434,6 +439,11 @@ export interface CommandDefinition<CommandName extends string> {
     previousCommandFields: ESQLFieldWithMetadata[],
     userDefinedColumns: ESQLFieldWithMetadata[]
   ) => ESQLFieldWithMetadata[];
+
+  /**
+   * This method is called to define the commands available after this command is applied.
+   */
+  commandsSuggestionsAfter?: (suggestions: SuggestionRawDefinition[]) => SuggestionRawDefinition[];
 }
 
 export interface CommandTypeDefinition {