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

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';
@@ -590,6 +592,16 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
   },
   {
     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,
     name: 'change_point',
     preview: true,
     description: i18n.translate(

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';
@@ -415,7 +416,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.