[ES|QL] support SAMPLE command

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

@@ -0,0 +1,96 @@
+/*
+ * 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 { Parser } from '../parser';
+
+describe('SAMPLE', () => {
+  describe('correctly formatted', () => {
+    test('without seed', () => {
+      const text = `
+        FROM employees
+        | SAMPLE 0.25
+        `;
+      const { ast, errors } = Parser.parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'sample',
+          args: [
+            {
+              type: 'literal',
+              literalType: 'double',
+              value: 0.25,
+            },
+          ],
+        },
+      ]);
+    });
+
+    test('with seed', () => {
+      const text = `
+        FROM employees
+        | SAMPLE 0.25 123
+        `;
+      const { ast, errors } = Parser.parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'sample',
+          args: [
+            {
+              type: 'literal',
+              literalType: 'double',
+              value: 0.25,
+            },
+            {
+              type: 'literal',
+              literalType: 'integer',
+              value: 123,
+            },
+          ],
+        },
+      ]);
+    });
+  });
+
+  describe('errors', () => {
+    it('wrong data type for probability', () => {
+      const { errors } = Parser.parse(`
+        FROM employees
+        | SAMPLE 25
+        `);
+
+      expect(errors.length).toBe(1);
+    });
+
+    it('wrong data type for seed', () => {
+      const { errors } = Parser.parse(`
+        FROM employees
+        | SAMPLE .25 .123
+        `);
+
+      expect(errors.length).toBe(1);
+    });
+
+    it('command with no args', () => {
+      const { errors } = Parser.parse(`
+        FROM employees
+        | SAMPLE 
+        `);
+
+      expect(errors.length).toBe(1);
+    });
+  });
+});

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

@@ -33,10 +33,11 @@ import {
   type WhereCommandContext,
   RerankCommandContext,
   RrfCommandContext,
+  SampleCommandContext,
 } from '../antlr/esql_parser';
 import { default as ESQLParserListener } from '../antlr/esql_parser_listener';
 import type { ESQLAst } from '../types';
-import { createCommand, createFunction, textExistsAndIsValid } from './factories';
+import { createCommand, createFunction, createLiteral, textExistsAndIsValid } from './factories';
 import { createChangePointCommand } from './factories/change_point';
 import { createDissectCommand } from './factories/dissect';
 import { createEvalCommand } from './factories/eval';
@@ -352,6 +353,18 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
     this.ast.push(command);
   }
 
+  exitSampleCommand(ctx: SampleCommandContext): void {
+    const command = createCommand('sample', ctx);
+    this.ast.push(command);
+
+    if (ctx._probability) {
+      command.args.push(createLiteral('double', ctx._probability.DECIMAL_LITERAL()));
+    }
+    if (ctx._seed) {
+      command.args.push(createLiteral('integer', ctx._seed.INTEGER_LITERAL()));
+    }
+  }
+
   /**
    * Exit a parse tree produced by `esql_parser.rrfCommand`.
    *

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

@@ -236,6 +236,22 @@ describe('single line query', () => {
       });
     });
 
+    describe('SAMPLE', () => {
+      test('from single line', () => {
+        const { text } = reprint(`FROM index | SAMPLE 0.1 123`);
+
+        expect(text).toBe('FROM index | SAMPLE 0.1 123');
+      });
+
+      test('from multiline', () => {
+        const { text } = reprint(`FROM index
+| SAMPLE 0.1 123
+          `);
+
+        expect(text).toBe('FROM index | SAMPLE 0.1 123');
+      });
+    });
+
     describe('RRF', () => {
       test('from single line', () => {
         const { text } =

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/constants.ts

@@ -25,7 +25,7 @@
  * DISSECT input "pattern"
  * ```
  */
-export const commandsWithNoCommaArgSeparator = new Set(['grok', 'dissect']);
+export const commandsWithNoCommaArgSeparator = new Set(['grok', 'dissect', 'sample']);
 
 /**
  * This set tracks command options which use an equals sign to separate

src/platform/packages/shared/kbn-esql-ast/src/synth/__tests__/cmd.test.ts

@@ -27,6 +27,13 @@ test('can create a RERANK command', () => {
   );
 });
 
+test('can create a SAMPLE command', () => {
+  const node = cmd`SAMPLE 0.23 123`;
+  const text = node.toString();
+
+  expect(text).toBe('SAMPLE 0.23 123');
+});
+
 test('can create a complex STATS command', () => {
   const node = cmd`STATS count_last_hour = SUM(count_last_hour), total_visits = SUM(total_visits), bytes_transform = SUM(bytes_transform), bytes_transform_last_hour = SUM(bytes_transform_last_hour) BY extension.keyword`;
   const text = BasicPrettyPrinter.command(node);

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> {}
 
+// SAMPLE <probability> [SEED <seed>]
+export class SampleCommandVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData
+> extends CommandVisitorContext<Methods, Data, ESQLAstCommand> {}
+
 // RRF
 export class RrfCommandVisitorContext<
   Methods extends VisitorMethods = VisitorMethods,

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 'sample': {
+        if (!this.methods.visitSampleCommand) break;
+        return this.visitSampleCommand(parent, commandNode, input as any);
+      }
       case 'rrf': {
         if (!this.methods.visitRrfCommand) break;
         return this.visitRrfCommand(parent, commandNode, input as any);
@@ -421,6 +425,15 @@ export class GlobalVisitorContext<
     return this.visitWithSpecificContext('visitForkCommand', context, input);
   }
 
+  public visitSampleCommand(
+    parent: contexts.VisitorContext | null,
+    node: ESQLAstCommand,
+    input: types.VisitorInput<Methods, 'visitSampleCommand'>
+  ): types.VisitorOutput<Methods, 'visitSampleCommand'> {
+    const context = new contexts.ForkCommandVisitorContext(this, node, parent);
+    return this.visitWithSpecificContext('visitSampleCommand', context, input);
+  }
+
   public visitRrfCommand(
     parent: contexts.VisitorContext | null,
     node: ESQLAstCommand,

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

@@ -188,6 +188,7 @@ export interface VisitorMethods<
     any
   >;
   visitForkCommand?: Visitor<contexts.ForkCommandVisitorContext<Visitors, Data>, any, any>;
+  visitSampleCommand?: Visitor<contexts.SampleCommandVisitorContext<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>;

src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts

@@ -8,6 +8,7 @@
  */
 
 import { parse } from '../parser';
+import { Parser } from '../parser/parser';
 import {
   ESQLColumn,
   ESQLCommand,
@@ -106,6 +107,21 @@ describe('structurally can walk all nodes', () => {
       expect(columns.map(({ name }) => name).sort()).toStrictEqual(['c', 'd']);
     });
 
+    test('can traverse SAMPLE command', () => {
+      const { root } = Parser.parse('FROM index | SAMPLE 0.25');
+      const commands: ESQLCommand[] = [];
+      const literals: ESQLLiteral[] = [];
+
+      walk(root, {
+        visitCommand: (cmd) => commands.push(cmd),
+        visitLiteral: (lit) => literals.push(lit),
+      });
+
+      expect(commands.map(({ name }) => name).sort()).toStrictEqual(['from', 'sample']);
+      expect(literals.length).toBe(1);
+      expect(literals[0].value).toBe(0.25);
+    });
+
     test('"visitAny" can capture command nodes', () => {
       const { ast } = parse('FROM index | STATS a = 123 | WHERE 123 | LIMIT 10');
       const commands: ESQLCommand[] = [];

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

@@ -0,0 +1,23 @@
+/*
+ * 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('SAMPLE <percentage> [<seed>]', () => {
+    test('suggests percentages', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions('from a | SAMPLE /', ['.1 ', '.01 ', '.001 ']);
+    });
+
+    test('suggests pipe after number', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions('from a | SAMPLE .48 /', ['| ']);
+    });
+  });
+});

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

@@ -32,7 +32,9 @@ import { getDateHistogramCompletionItem } from './commands/stats/util';
 import { getSafeInsertText, TIME_SYSTEM_PARAMS, TRIGGER_SUGGESTION_COMMAND } from './factories';
 import { getRecommendedQueries } from './recommended_queries/templates';
 
-const commandDefinitions = unmodifiedCommandDefinitions.filter(({ hidden }) => !hidden);
+const commandDefinitions = unmodifiedCommandDefinitions.filter(
+  ({ name, hidden }) => !hidden && name !== 'rrf'
+);
 
 const getRecommendedQueriesSuggestions = (fromCommand: string, timeField?: string) =>
   getRecommendedQueries({

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/sample/index.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 { CommandSuggestParams } from '../../../definitions/types';
+import type { SuggestionRawDefinition } from '../../types';
+import { buildConstantsDefinitions } from '../../factories';
+import { pipeCompleteItem } from '../../complete_items';
+
+export function suggest({ innerText }: CommandSuggestParams<'sample'>): SuggestionRawDefinition[] {
+  // test for a number and at least one whitespace char at the end of the innerText
+  if (/[0-9]\s+$/.test(innerText)) {
+    return [pipeCompleteItem];
+  }
+
+  return buildConstantsDefinitions(['.1', '.01', '.001'], '', undefined, {
+    advanceCursorAndOpenSuggestions: true,
+  });
+}

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

@@ -83,6 +83,7 @@ import {
   suggest as suggestForChangePoint,
   fieldsSuggestionsAfter as fieldsSuggestionsAfterChangePoint,
 } from '../autocomplete/commands/change_point';
+import { suggest as suggestForSample } from '../autocomplete/commands/sample';
 
 import { METADATA_FIELDS } from '../shared/constants';
 import { getMessageFromId } from '../validation/errors';
@@ -710,6 +711,17 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
   },
   {
     hidden: true,
+    name: 'sample',
+    preview: true,
+    description: i18n.translate('kbn-esql-validation-autocomplete.esql.definitions.sampleDoc', {
+      defaultMessage:
+        'Samples a percentage of the results, optionally with a seed for reproducibility.',
+    }),
+    declaration: `SAMPLE <percentage> [<seed>]`,
+    examples: [],
+    suggest: suggestForSample,
+  },
+  {
     preview: true,
     name: 'rrf',
     description: i18n.translate('kbn-esql-validation-autocomplete.esql.definitions.rrfDoc', {