[ES|QL] Add support for RRF (#221349)

## Summary

Part of #215092

## Considerations
### `RRF` may change in the future
@ioanatia has stated that the command may change in the future as it's
not fully implemented yet.

### Licence check
Will be addressed at #216791



https://github.com/user-attachments/assets/759470a2-4afa-4d20-adbf-cb3001895e95


## New command support checklist
### AST support
- [x] Make sure that the new command is in the local Kibana grammar
definition. The ANTLR lexer and parser files are updated every Monday
from the source definition of the language at Elasticsearch (via a
manually merged, automatically generated
[PR](#213006)).
- [x] Create a factory for generating the new node. The new node should
satisfy the `ESQLCommand<Name>` interface. If the syntax of your command
cannot be decomposed only in parameters, you can hold extra information
by extending the `ESQLCommand` interface. I.E., check the Rerank
command.
- [x] While ANTLR is parsing the text query, we create our own AST by
using `onExit` listeners at
`kbn-esql-ast/src/parser/esql_ast_builder_listener.ts`. Implement the
`onExit<COMMAND_NAME>` method based on the interface autogenerated by
ANTLR and push the new node into the AST.
- [x] Create unit tests checking that the correct AST nodes are
generated when parsing your command.
- [x] Add a dedicated [visitor
callback](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-esql-ast/src/visitor/README.md)
for the new command.
- [x] Verify that the `Walker` API can visit the new node.
- [x] Verify that the `Synth` API can construct the new node.

### Validating that the command works well when prettifying the query
- [x] Validate that the prettifier works correctly.
- [ ] Adjust the basic pretty printer and the wrapping pretty printer if
needed.
- [x] Add unit tests validating that the queries are correctly formatted
(even if no adjustment has been done).


### Creating the command definition
- [x] Add the definition of the new command at
`kbn-esql-validation-autocomplete/src/definitions/commands.ts`.

### Adding the corresponding client-side validations
- [x] Add a custom validation if needed.
- [x] Add tests checking the behavior of the validation following this
[guide](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-esql-validation-autocomplete/README.md#the-new-way).


### Adding the autocomplete suggestions
- [x] Add the suggestions to be shown when **positioned at** the new
command.
- [x] Create a new folder at
`kbn-esql-validation-autocomplete/src/autocomplete/commands` for your
command.
- [x] Export a `suggest` function that should return an array of
suggestions and set it up into the command definition.
- [ ] Optionally, we must filter or incorporate fields suggestions after
a command is executed, this is supported by adding the
`fieldsSuggestionsAfter` method to the command definition. Read this
documentation to understand how it works.
- [x] If the new command must be suggested only in particular
situations, modify the corresponding suggestions to make it possible.
- [x] Add tests following this
[guide](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-esql-validation-autocomplete/README.md#automated-testing).


### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [ ] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...

---------

Co-authored-by: Elastic Machine <[email protected]>
Co-authored-by: Stratoula Kalafateli <[email protected]>
Co-authored-by: Drew Tate <[email protected]>

Author: 4 people

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

@@ -0,0 +1,68 @@
+/*
+ * 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: [],
+      });
+    });
+  });
+
+  describe('when incorrectly formatted, return errors', () => {
+    it('when no pipe after', () => {
+      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 { errors } = parse(text);
+
+      expect(errors.length > 0).toBe(true);
+    });
+
+    it('when no pipe between FORK and RRF', () => {
+      const text = `FROM search-movies METADATA _score, _id, _index
+                    | FORK
+                      ( WHERE semantic_title:"Shakespeare" | SORT _score)
+                      ( WHERE title:"Shakespeare" | SORT _score) RRF`;
+
+      const { errors } = parse(text);
+
+      expect(errors.length > 0).toBe(true);
+    });
+
+    it('when RRF is invoked with arguments', () => {
+      const text = `FROM search-movies METADATA _score, _id, _index
+                    | FORK ( WHERE semantic_title:"Shakespeare" | SORT _score)
+                              ( WHERE title:"Shakespeare" | SORT _score)
+                    | RRF text`;
+
+      const { errors } = parse(text);
+
+      expect(errors.length > 0).toBe(true);
+    });
+  });
+});

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

@@ -12,6 +12,12 @@ import type { IndexAutocompleteItem } from '@kbn/esql-types';
 import { ESQLFieldWithMetadata } from '../validation/types';
 import { fieldTypes } from '../definitions/types';
 import { ESQLCallbacks } from '../shared/types';
+import { METADATA_FIELDS } from '../shared/constants';
+
+export const metadataFields: ESQLFieldWithMetadata[] = METADATA_FIELDS.map((field) => ({
+  name: field,
+  type: 'keyword',
+}));
 
 export const fields: ESQLFieldWithMetadata[] = [
   ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type })),
@@ -107,6 +113,9 @@ export function getCallbackMocks(): ESQLCallbacks {
         };
         return [field];
       }
+      if (/METADATA/i.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,53 @@
+/*
+ * 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';
+
+jest.mock('../../definitions/commands', () => {
+  const actual = jest.requireActual('../../definitions/commands');
+  const modifiedDefinitions = actual.commandDefinitions.map((def: any) =>
+    def.name === 'rrf' ? { ...def, hidden: false } : def
+  );
+  return {
+    ...actual,
+    commandDefinitions: modifiedDefinitions,
+  };
+});
+
+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 ');
+      expect(suggestedCommands).toContain('LIMIT ');
+    });
+
+    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,12 @@ export async function suggest(
       return [...sourceCommandsSuggestions, ...recommendedQueriesSuggestions];
     }
 
+    // If the last command is not a FORK, RRF should not be suggested.
+    const lastCommand = root.commands[root.commands.length - 1];
+    if (lastCommand.name !== 'fork') {
+      suggestions = suggestions.filter((def) => def.label !== 'RRF');
+    }
+
     return suggestions.filter((def) => !isSourceCommand(def));
   }
 

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

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

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

@@ -69,6 +69,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';
@@ -706,4 +708,17 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
 
     fieldsSuggestionsAfter: fieldsSuggestionsAfterFork,
   },
+  {
+    hidden: true,
+    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`,
+    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.