diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/sample.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/sample.test.ts new file mode 100644 index 0000000000000..7df9c4a4e1555 --- /dev/null +++ b/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); + }); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts index aad759c5c4a90..672696edca344 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts +++ b/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`. * diff --git a/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts index 83dae2c4821c7..2049c5318a2d0 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts +++ b/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 } = diff --git a/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/constants.ts b/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/constants.ts index 01208af98d025..f2076d2f4e168 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/pretty_print/constants.ts +++ b/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 diff --git a/src/platform/packages/shared/kbn-esql-ast/src/synth/__tests__/cmd.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/synth/__tests__/cmd.test.ts index 91b05de857c40..34fd68b22ecd4 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/synth/__tests__/cmd.test.ts +++ b/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); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts index 721d11a04dfba..332c41ee6b68e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts @@ -520,6 +520,12 @@ export class ForkCommandVisitorContext< Data extends SharedData = SharedData > extends CommandVisitorContext {} +// SAMPLE [SEED ] +export class SampleCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + // RRF export class RrfCommandVisitorContext< Methods extends VisitorMethods = VisitorMethods, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts index 45c3f3619ed1a..15bd2e02a4acf 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts +++ b/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 + ): types.VisitorOutput { + const context = new contexts.ForkCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitSampleCommand', context, input); + } + public visitRrfCommand( parent: contexts.VisitorContext | null, node: ESQLAstCommand, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts index 25e5df0bda2ec..dc29f0faf28c3 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts @@ -188,6 +188,7 @@ export interface VisitorMethods< any >; visitForkCommand?: Visitor, any, any>; + visitSampleCommand?: Visitor, any, any>; visitCommandOption?: Visitor, any, any>; visitRrfCommand?: Visitor, any, any>; visitExpression?: Visitor, any, any>; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts index cf3b2670d3ed5..8ad5ec2070d83 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts +++ b/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[] = []; diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.sample.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.sample.test.ts new file mode 100644 index 0000000000000..8e4f034d68991 --- /dev/null +++ b/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 []', () => { + 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 /', ['| ']); + }); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts index 7c22b2214bb73..31ef0badb9b8c 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts +++ b/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({ diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/sample/index.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/sample/index.ts new file mode 100644 index 0000000000000..dca685761ab2b --- /dev/null +++ b/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, + }); +} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/definitions/commands.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/definitions/commands.ts index a28b2c2354f2e..cb4ad02d9c8f8 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/definitions/commands.ts +++ b/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> = [ }, { 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 []`, + examples: [], + suggest: suggestForSample, + }, + { preview: true, name: 'rrf', description: i18n.translate('kbn-esql-validation-autocomplete.esql.definitions.rrfDoc', {