[ES|QL] STATS command field WHERE expression support (#206593)

## Summary

Closes https://github.com/elastic/kibana/issues/195363

Adds support for ` ... WHERE ... ` expressions inside `STATS` command.
Like

```
FROM index | STATS a WHERE b
```


### Checklist

- [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

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Co-authored-by: Stratoula Kalafateli <efstratia.kalafateli@elastic.co>

src/platform/packages/shared/kbn-esql-ast/index.ts

@@ -31,10 +31,12 @@ export type {
 } from './src/types';
 
 export {
-  isBinaryExpression,
   isColumn,
   isDoubleLiteral,
   isFunctionExpression,
+  isBinaryExpression,
+  isWhereExpression,
+  isFieldExpression,
   isIdentifier,
   isIntegerLiteral,
   isLiteral,

src/platform/packages/shared/kbn-esql-ast/src/ast/helpers.ts

@@ -8,6 +8,7 @@
  */
 
 import type {
+  BinaryExpressionWhereOperator,
   ESQLAstNode,
   ESQLBinaryExpression,
   ESQLColumn,
@@ -41,6 +42,16 @@ export const isFunctionExpression = (node: unknown): node is ESQLFunction =>
 export const isBinaryExpression = (node: unknown): node is ESQLBinaryExpression =>
   isFunctionExpression(node) && node.subtype === 'binary-expression';
 
+export const isWhereExpression = (
+  node: unknown
+): node is ESQLBinaryExpression<BinaryExpressionWhereOperator> =>
+  isBinaryExpression(node) && node.name === 'where';
+
+export const isFieldExpression = (
+  node: unknown
+): node is ESQLBinaryExpression<BinaryExpressionWhereOperator> =>
+  isBinaryExpression(node) && node.name === '=';
+
 export const isLiteral = (node: unknown): node is ESQLLiteral =>
   isProperNode(node) && node.type === 'literal';
 

src/platform/packages/shared/kbn-esql-ast/src/builder/builder.ts

@@ -38,6 +38,7 @@ import {
   ESQLTimeInterval,
   ESQLBooleanLiteral,
   ESQLNullLiteral,
+  BinaryExpressionOperator,
 } from '../types';
 import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types';
 
@@ -56,10 +57,10 @@ export namespace Builder {
     incomplete,
   });
 
-  export const command = (
-    template: PartialFields<AstNodeTemplate<ESQLCommand>, 'args'>,
+  export const command = <Name extends string>(
+    template: PartialFields<AstNodeTemplate<ESQLCommand<Name>>, 'args'>,
     fromParser?: Partial<AstNodeParserFields>
-  ): ESQLCommand => {
+  ): ESQLCommand<Name> => {
     return {
       ...template,
       ...Builder.parserFields(fromParser),
@@ -260,20 +261,26 @@ export namespace Builder {
         ) as ESQLUnaryExpression;
       };
 
-      export const binary = (
-        name: string,
+      export const binary = <Name extends BinaryExpressionOperator = BinaryExpressionOperator>(
+        name: Name,
         args: [left: ESQLAstItem, right: ESQLAstItem],
         template?: Omit<AstNodeTemplate<ESQLFunction>, 'subtype' | 'name' | 'operator' | 'args'>,
         fromParser?: Partial<AstNodeParserFields>
-      ): ESQLBinaryExpression => {
+      ): ESQLBinaryExpression<Name> => {
         const operator = Builder.identifier({ name });
         return Builder.expression.func.node(
           { ...template, name, operator, args, subtype: 'binary-expression' },
           fromParser
-        ) as ESQLBinaryExpression;
+        ) as ESQLBinaryExpression<Name>;
       };
     }
 
+    export const where = (
+      args: [left: ESQLAstItem, right: ESQLAstItem],
+      template?: Omit<AstNodeTemplate<ESQLFunction>, 'subtype' | 'name' | 'operator' | 'args'>,
+      fromParser?: Partial<AstNodeParserFields>
+    ) => Builder.expression.func.binary('where', args, template, fromParser);
+
     export namespace literal {
       /**
        * Constructs a NULL literal node.

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

@@ -0,0 +1,319 @@
+/*
+ * 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 { EsqlQuery } from '../../query';
+import { Walker } from '../../walker';
+
+describe('STATS', () => {
+  describe('correctly formatted', () => {
+    it('a simple single aggregate expression', () => {
+      const src = `
+        FROM employees
+          | STATS 123
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'literal',
+              name: '123',
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('aggregation function with escaped values', () => {
+      const src = `
+        FROM employees
+          | STATS 123, agg("salary")
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'literal',
+              value: 123,
+            },
+            {
+              type: 'function',
+              name: 'agg',
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('field column name defined', () => {
+      const src = `
+        FROM employees
+          | STATS my_field = agg("salary")
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'function',
+              name: '=',
+              args: [
+                {
+                  type: 'column',
+                  args: [
+                    {
+                      type: 'identifier',
+                      name: 'my_field',
+                    },
+                  ],
+                },
+                [
+                  {
+                    type: 'function',
+                    name: 'agg',
+                    args: [
+                      {
+                        type: 'literal',
+                        valueUnquoted: 'salary',
+                      },
+                    ],
+                  },
+                ],
+              ],
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('parses BY clause', () => {
+      const src = `
+        FROM employees
+          | STATS my_field = agg("salary") BY department
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {},
+            {
+              type: 'option',
+              name: 'by',
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    describe('WHERE clause', () => {
+      it('boolean expression wrapped in WHERE clause', () => {
+        const src = `
+          FROM employees
+            | STATS 123 WHERE still_hired == true
+            | LIMIT 1
+            `;
+        const query = EsqlQuery.fromSrc(src);
+
+        // console.log(JSON.stringify(query.ast.commands, null, 2));
+
+        expect(query.errors.length).toBe(0);
+        expect(query.ast.commands).toMatchObject([
+          {},
+          {
+            type: 'command',
+            name: 'stats',
+            args: [
+              {
+                type: 'function',
+                subtype: 'binary-expression',
+                name: 'where',
+                args: [
+                  {
+                    type: 'literal',
+                    name: '123',
+                  },
+                  {
+                    type: 'function',
+                    name: '==',
+                    args: [
+                      {
+                        type: 'column',
+                        args: [
+                          {
+                            type: 'identifier',
+                            name: 'still_hired',
+                          },
+                        ],
+                      },
+                      {
+                        type: 'literal',
+                        value: 'true',
+                      },
+                    ],
+                  },
+                ],
+              },
+            ],
+          },
+          {},
+        ]);
+      });
+
+      it('extracts WHERE position', () => {
+        const src = `
+          FROM employees
+            | STATS 123 WHERE still_hired == true
+            | LIMIT 1
+            `;
+        const query = EsqlQuery.fromSrc(src);
+        const where = Walker.match(query.ast, {
+          type: 'function',
+          name: 'where',
+        })!;
+        const text = src.substring(where.location.min, where.location.max + 1);
+
+        expect(text.trim()).toBe('123 WHERE still_hired == true');
+      });
+
+      it('WHERE clause around "agg" function', () => {
+        const src = `
+        FROM employees
+          | STATS 123, agg("salary") WHERE 456, 789 
+          | LIMIT 10
+          `;
+        const query = EsqlQuery.fromSrc(src);
+
+        expect(query.errors.length).toBe(0);
+        expect(query.ast.commands).toMatchObject([
+          {},
+          {
+            type: 'command',
+            name: 'stats',
+            args: [
+              {
+                type: 'literal',
+                value: 123,
+              },
+              {
+                type: 'function',
+                name: 'where',
+                args: [
+                  {
+                    type: 'function',
+                    name: 'agg',
+                  },
+                  {
+                    type: 'literal',
+                    value: 456,
+                  },
+                ],
+              },
+              {
+                type: 'literal',
+                value: 789,
+              },
+            ],
+          },
+          {},
+        ]);
+      });
+
+      it('WHERE for field definition', () => {
+        const src = `
+          FROM employees
+            | STATS my_field = agg("salary") WHERE 123, 456
+            | WHERE still_hired == true
+            `;
+        const query = EsqlQuery.fromSrc(src);
+
+        expect(query.errors.length).toBe(0);
+        expect(query.ast.commands).toMatchObject([
+          {},
+          {
+            type: 'command',
+            name: 'stats',
+            args: [
+              {
+                type: 'function',
+                name: 'where',
+                args: [
+                  {
+                    type: 'function',
+                    name: '=',
+                    args: [
+                      {
+                        type: 'column',
+                        args: [
+                          {
+                            type: 'identifier',
+                            name: 'my_field',
+                          },
+                        ],
+                      },
+                      [
+                        {
+                          type: 'function',
+                          name: 'agg',
+                          args: [
+                            {
+                              type: 'literal',
+                              valueUnquoted: 'salary',
+                            },
+                          ],
+                        },
+                      ],
+                    ],
+                  },
+                  {
+                    type: 'literal',
+                    value: 123,
+                  },
+                ],
+              },
+              {
+                type: 'literal',
+                value: 456,
+              },
+            ],
+          },
+          {},
+        ]);
+      });
+    });
+  });
+});

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

@@ -60,10 +60,13 @@ import type { ESQLAst, ESQLAstMetricsCommand } from '../types';
 import { createJoinCommand } from './factories/join';
 import { createDissectCommand } from './factories/dissect';
 import { createGrokCommand } from './factories/grok';
+import { createStatsCommand } from './factories/stats';
 
 export class ESQLAstBuilderListener implements ESQLParserListener {
   private ast: ESQLAst = [];
 
+  constructor(public src: string) {}
+
   public getAst() {
     return { ast: this.ast };
   }
@@ -171,16 +174,9 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitStatsCommand(ctx: StatsCommandContext) {
-    const command = createCommand('stats', ctx);
-    this.ast.push(command);
+    const command = createStatsCommand(ctx, this.src);
 
-    // STATS expression is optional
-    if (ctx._stats) {
-      command.args.push(...collectAllAggFields(ctx.aggFields()));
-    }
-    if (ctx._grouping) {
-      command.args.push(...visitByOption(ctx, ctx.fields()));
-    }
+    this.ast.push(command);
   }
 
   /**

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

@@ -84,7 +84,7 @@ const createParserFields = (ctx: ParserRuleContext): AstNodeParserFields => ({
   incomplete: Boolean(ctx.exception),
 });
 
-export const createCommand = (name: string, ctx: ParserRuleContext) =>
+export const createCommand = <Name extends string>(name: Name, ctx: ParserRuleContext) =>
   Builder.command({ name, args: [] }, createParserFields(ctx));
 
 export const createInlineCast = (ctx: InlineCastContext, value: ESQLInlineCast['value']) =>

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

@@ -0,0 +1,64 @@
+/*
+ * 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 { AggFieldContext, FieldContext, StatsCommandContext } from '../../antlr/esql_parser';
+import { Builder } from '../../builder';
+import { ESQLCommand } from '../../types';
+import { firstItem, resolveItem } from '../../visitor/utils';
+import { createCommand } from '../factories';
+import { collectBooleanExpression, visitByOption, visitField } from '../walkers';
+
+const createField = (ctx: FieldContext) => visitField(ctx)[0];
+
+const createAggField = (ctx: AggFieldContext) => {
+  const fieldCtx = ctx.field();
+  const field = createField(fieldCtx);
+
+  const booleanExpression = ctx.booleanExpression();
+
+  if (!booleanExpression) {
+    return field;
+  }
+
+  const condition = collectBooleanExpression(booleanExpression)[0];
+  const aggField = Builder.expression.where(
+    [field, condition],
+    {},
+    {
+      location: {
+        min: firstItem([resolveItem(field)])?.location?.min ?? 0,
+        max: firstItem([resolveItem(condition)])?.location?.max ?? 0,
+      },
+    }
+  );
+
+  return aggField;
+};
+
+export const createStatsCommand = (ctx: StatsCommandContext, src: string): ESQLCommand<'stats'> => {
+  const command = createCommand('stats', ctx);
+
+  if (ctx._stats) {
+    const fields = ctx.aggFields();
+
+    for (const fieldCtx of fields.aggField_list()) {
+      const node = createAggField(fieldCtx);
+
+      command.args.push(node);
+    }
+  }
+
+  if (ctx._grouping) {
+    const options = visitByOption(ctx, ctx.fields());
+
+    command.args.push(...options);
+  }
+
+  return command;
+};

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

@@ -55,11 +55,11 @@ export const getParser = (
   };
 };
 
-export const createParser = (text: string) => {
+export const createParser = (src: string) => {
   const errorListener = new ESQLErrorListener();
-  const parseListener = new ESQLAstBuilderListener();
+  const parseListener = new ESQLAstBuilderListener(src);
 
-  return getParser(CharStreams.fromString(text), errorListener, parseListener);
+  return getParser(CharStreams.fromString(src), errorListener, parseListener);
 };
 
 // These will need to be manually updated whenever the relevant grammar changes.
@@ -106,7 +106,7 @@ export const parse = (text: string | undefined, options: ParseOptions = {}): Par
       return { ast: commands, root: Builder.expression.query(commands), errors: [], tokens: [] };
     }
     const errorListener = new ESQLErrorListener();
-    const parseListener = new ESQLAstBuilderListener();
+    const parseListener = new ESQLAstBuilderListener(text);
     const { tokens, parser } = getParser(
       CharStreams.fromString(text),
       errorListener,

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

@@ -367,6 +367,18 @@ describe('single line query', () => {
 
           expect(text).toBe('FROM a | WHERE (1 + 1 + 2) * ((3 - 4) / (5 + 6 + 7) + 1)');
         });
+
+        test('formats WHERE binary-expression', () => {
+          const { text } = reprint('FROM a | STATS a WHERE b');
+
+          expect(text).toBe('FROM a | STATS a WHERE b');
+        });
+
+        test('formats complex WHERE binary-expression', () => {
+          const { text } = reprint('FROM a | STATS a = agg(123) WHERE b == test(c, 123)');
+
+          expect(text).toBe('FROM a | STATS a = AGG(123) WHERE b == TEST(c, 123)');
+        });
       });
     });
 

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

@@ -583,6 +583,15 @@ ROW
     // Two is more important here
     2`);
       });
+
+      test('WHERE expression', () => {
+        const query = `FROM a | STATS /* 1 */ a /* 2 */ WHERE /* 3 */ a /* 4 */ == /* 5 */ 1 /* 6 */`;
+        const text = reprint(query).text;
+
+        expect(text).toBe(
+          'FROM a | STATS /* 1 */ a /* 2 */ WHERE /* 3 */ a /* 4 */ == /* 5 */ 1 /* 6 */'
+        );
+      });
     });
   });
 });

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

@@ -657,6 +657,22 @@ FROM index
           11111111111111.111 +
           11111111111111.111)`);
       });
+
+      test('formats WHERE binary-expression', () => {
+        const query = `
+        FROM index
+        | STATS fn(11111111111111 - 11111111111111 - 11111111111111 - 11111111111111) WHERE 11111111111111 == AHA(11111111111111 + 11111111111111 + 11111111111111),
+        | LIMIT 10
+        `;
+        const text = reprint(query).text;
+
+        expect('\n' + text).toBe(`
+FROM index
+  | STATS
+      FN(11111111111111 - 11111111111111 - 11111111111111 - 11111111111111) WHERE
+        11111111111111 == AHA(11111111111111 + 11111111111111 + 11111111111111)
+  | LIMIT 10`);
+      });
     });
   });
 

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

@@ -322,6 +322,7 @@ export class WrappingPrettyPrinter {
       lines = 1;
       txt = ctx instanceof CommandVisitorContext ? '' : '\n';
       const args = [...ctx.arguments()].filter((arg) => {
+        if (!arg) return false;
         if (arg.type === 'option') return arg.name === 'as';
         return true;
       });
@@ -397,6 +398,7 @@ export class WrappingPrettyPrinter {
             indented = true;
           }
           txt = indent + LeafPrinter.comment(decoration) + '\n' + txt;
+          indented = true;
         }
       }
     }
@@ -476,17 +478,23 @@ export class WrappingPrettyPrinter {
         !(value.type === 'function' && value.subtype === 'variadic-call');
       const castType = ctx.node.castType;
 
-      let valueFormatted = ctx.visitValue({
+      const valueResult = ctx.visitValue({
         indent: inp.indent,
         remaining: inp.remaining - castType.length - 2,
-      }).txt;
+      });
+      let { txt: valueFormatted } = valueResult;
 
       if (wrapInBrackets) {
         valueFormatted = `(${valueFormatted})`;
       }
 
       const formatted = `${valueFormatted}::${ctx.node.castType}${inp.suffix ?? ''}`;
-      const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted);
+      const { txt, indented } = this.decorateWithComments(
+        inp.indent,
+        ctx.node,
+        formatted,
+        valueResult.indented
+      );
 
       return { txt, indented };
     })

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

@@ -193,8 +193,9 @@ export interface ESQLOrderExpression extends ESQLAstBaseItem {
   args: [field: ESQLAstItem];
 }
 
-export interface ESQLBinaryExpression
-  extends ESQLFunction<'binary-expression', BinaryExpressionOperator> {
+export interface ESQLBinaryExpression<
+  Name extends BinaryExpressionOperator = BinaryExpressionOperator
+> extends ESQLFunction<'binary-expression', Name> {
   subtype: 'binary-expression';
   args: [ESQLAstItem, ESQLAstItem];
 }
@@ -204,13 +205,15 @@ export type BinaryExpressionOperator =
   | BinaryExpressionAssignmentOperator
   | BinaryExpressionComparisonOperator
   | BinaryExpressionRegexOperator
-  | BinaryExpressionRenameOperator;
+  | BinaryExpressionRenameOperator
+  | BinaryExpressionWhereOperator;
 
 export type BinaryExpressionArithmeticOperator = '+' | '-' | '*' | '/' | '%';
 export type BinaryExpressionAssignmentOperator = '=';
 export type BinaryExpressionComparisonOperator = '==' | '=~' | '!=' | '<' | '<=' | '>' | '>=';
 export type BinaryExpressionRegexOperator = 'like' | 'not_like' | 'rlike' | 'not_rlike';
 export type BinaryExpressionRenameOperator = 'as';
+export type BinaryExpressionWhereOperator = 'where';
 
 // from https://github.com/elastic/elasticsearch/blob/122e7288200ee03e9087c98dff6cebbc94e774aa/docs/reference/esql/functions/kibana/inline_cast.json
 export type InlineCastingType =

src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts

@@ -159,6 +159,34 @@ test('"visitLiteral" takes over all literal visits', () => {
   expect(text).toBe('FROM E | STATS <LITERAL>, <LITERAL>, E, E | LIMIT <LITERAL>');
 });
 
+test('"visitExpression" does visit WHERE clause args', () => {
+  const { ast } = parse(`
+    FROM index
+      | STATS 1 WHERE 2
+      | LIMIT 123
+  `);
+  const visitor = new Visitor()
+    .on('visitLiteralExpression', (ctx) => {
+      return '<LITERAL>';
+    })
+    .on('visitFunctionCallExpression', (ctx) => {
+      return `${ctx.node.name}(${[...ctx.visitArguments(undefined)].join(', ')})`;
+    })
+    .on('visitExpression', (ctx) => {
+      return 'E';
+    })
+    .on('visitCommand', (ctx) => {
+      const args = [...ctx.visitArguments()].join(', ');
+      return `${ctx.name()}${args ? ` ${args}` : ''}`;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()].join(' | ');
+    });
+  const text = visitor.visitQuery(ast);
+
+  expect(text).toBe('FROM E | STATS where(<LITERAL>, <LITERAL>) | LIMIT <LITERAL>');
+});
+
 test('"visitExpression" does visit identifier nodes', () => {
   const { ast } = parse(`
     FROM index

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

@@ -18,8 +18,6 @@ import type {
   ESQLAstExpression,
   ESQLAstItem,
   ESQLAstJoinCommand,
-  ESQLAstNodeWithArgs,
-  ESQLAstNodeWithChildren,
   ESQLAstRenameExpression,
   ESQLColumn,
   ESQLCommandOption,
@@ -46,14 +44,7 @@ import type {
   VisitorOutput,
 } from './types';
 import { Builder } from '../builder';
-
-const isNodeWithArgs = (x: unknown): x is ESQLAstNodeWithArgs =>
-  !!x && typeof x === 'object' && Array.isArray((x as any).args);
-
-const isNodeWithChildren = (x: unknown): x is ESQLAstNodeWithChildren =>
-  !!x &&
-  typeof x === 'object' &&
-  (Array.isArray((x as any).args) || Array.isArray((x as any).values));
+import { isProperNode } from '../ast/helpers';
 
 export class VisitorContext<
   Methods extends VisitorMethods = VisitorMethods,
@@ -85,13 +76,8 @@ export class VisitorContext<
   ): Iterable<VisitorOutput<Methods, 'visitExpression'>> {
     this.ctx.assertMethodExists('visitExpression');
 
-    const node = this.node;
-
-    if (!isNodeWithArgs(node)) {
-      return;
-    }
-
-    for (const arg of singleItems(node.args)) {
+    for (const arg of this.arguments()) {
+      if (!arg) continue;
       if (arg.type === 'option' && arg.name !== 'as') {
         continue;
       }
@@ -107,7 +93,7 @@ export class VisitorContext<
   public arguments(): ESQLAstExpressionNode[] {
     const node = this.node;
 
-    if (!isNodeWithChildren(node)) {
+    if (!isProperNode(node)) {
       return [];
     }
 
@@ -128,12 +114,12 @@ export class VisitorContext<
 
     const node = this.node;
 
-    if (!isNodeWithArgs(node)) {
+    if (!isProperNode(node)) {
       throw new Error('Node does not have arguments');
     }
 
     let i = 0;
-    for (const arg of singleItems(node.args)) {
+    for (const arg of this.arguments()) {
       if (i === index) {
         return this.visitExpression(arg, input as any);
       }
@@ -193,7 +179,7 @@ export class CommandVisitorContext<
 
   public *options(): Iterable<ESQLCommandOption> {
     for (const arg of this.node.args) {
-      if (Array.isArray(arg)) {
+      if (!arg || Array.isArray(arg)) {
         continue;
       }
       if (arg.type === 'option') {
@@ -222,6 +208,9 @@ export class CommandVisitorContext<
 
     if (!option) {
       for (const arg of this.node.args) {
+        if (!arg) {
+          continue;
+        }
         if (Array.isArray(arg)) {
           yield arg;
           continue;
@@ -236,7 +225,7 @@ export class CommandVisitorContext<
     }
 
     const optionNode = this.node.args.find(
-      (arg) => !Array.isArray(arg) && arg.type === 'option' && arg.name === option
+      (arg) => !Array.isArray(arg) && arg && arg.type === 'option' && arg.name === option
     );
 
     if (optionNode) {

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

@@ -59,6 +59,7 @@ export function* children(node: ESQLProperNode): Iterable<ESQLAstExpression> {
   switch (node.type) {
     case 'function':
     case 'command':
+    case 'order':
     case 'option': {
       for (const arg of singleItems(node.args)) {
         yield arg;

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

@@ -218,6 +218,37 @@ describe('structurally can walk all nodes', () => {
             'index4',
           ]);
         });
+
+        test('can walk through "WHERE" binary expression', () => {
+          const query = 'FROM index | STATS a = 123 WHERE c == d';
+          const { root } = parse(query);
+          const expressions: ESQLFunction[] = [];
+
+          walk(root, {
+            visitFunction: (node) => {
+              if (node.name === 'where') {
+                expressions.push(node);
+              }
+            },
+          });
+
+          expect(expressions.length).toBe(1);
+          expect(expressions[0]).toMatchObject({
+            type: 'function',
+            subtype: 'binary-expression',
+            name: 'where',
+            args: [
+              {
+                type: 'function',
+                name: '=',
+              },
+              {
+                type: 'function',
+                name: '==',
+              },
+            ],
+          });
+        });
       });
 
       describe('columns', () => {
@@ -1189,11 +1220,11 @@ describe('Walker.matchAll()', () => {
 });
 
 describe('Walker.hasFunction()', () => {
-  test('can find assignment expression', () => {
+  test('can find binary expression expression', () => {
     const query1 = 'FROM a | STATS bucket(bytes, 1 hour)';
-    const query2 = 'FROM b | STATS var0 = bucket(bytes, 1 hour)';
-    const has1 = Walker.hasFunction(parse(query1).ast!, '=');
-    const has2 = Walker.hasFunction(parse(query2).ast!, '=');
+    const query2 = 'FROM b | STATS var0 == bucket(bytes, 1 hour)';
+    const has1 = Walker.hasFunction(parse(query1).ast!, '==');
+    const has2 = Walker.hasFunction(parse(query2).ast!, '==');
 
     expect(has1).toBe(false);
     expect(has2).toBe(true);

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

@@ -65,20 +65,6 @@ export type WalkerAstNode = ESQLAstNode | ESQLAstNode[];
 /**
  * Iterates over all nodes in the AST and calls the appropriate visitor
  * functions.
- *
- * AST nodes supported:
- *
- * - [x] command
- * - [x] option
- * - [x] mode
- * - [x] function
- * - [x] source
- * - [x] column
- * - [x] literal
- * - [x] list literal
- * - [x] timeInterval
- * - [x] inlineCast
- * - [x] unknown
  */
 export class Walker {
   /**
@@ -325,7 +311,7 @@ export class Walker {
     }
   }
 
-  public walkAstItem(node: ESQLAstItem): void {
+  public walkAstItem(node: ESQLAstItem | ESQLAstExpression): void {
     if (node instanceof Array) {
       const list = node as ESQLAstItem[];
       for (const item of list) this.walkAstItem(item);
@@ -373,7 +359,7 @@ export class Walker {
     const args = node.args;
     const length = args.length;
 
-    if (node.operator) this.walkAstItem(node.operator);
+    if (node.operator) this.walkSingleAstItem(node.operator);
 
     for (let i = 0; i < length; i++) {
       const arg = args[i];
@@ -393,6 +379,7 @@ export class Walker {
   }
 
   public walkSingleAstItem(node: ESQLAstExpression): void {
+    if (!node) return;
     const { options } = this;
     options.visitSingleAstItem?.(node);
     switch (node.type) {

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

@@ -645,6 +645,24 @@ const otherDefinitions: FunctionDefinition[] = [
       },
     ],
   },
+  {
+    type: 'builtin' as const,
+    name: 'where',
+    description: i18n.translate('kbn-esql-validation-autocomplete.esql.definition.whereDoc', {
+      defaultMessage: 'WHERE operator',
+    }),
+    supportedCommands: ['stats', 'inlinestats', 'metrics'],
+    supportedOptions: [],
+    signatures: [
+      {
+        params: [
+          { name: 'left', type: 'any' },
+          { name: 'right', type: 'any' },
+        ],
+        returnType: 'unknown',
+      },
+    ],
+  },
   {
     // TODO — this shouldn't be a function or an operator...
     name: 'info',

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

@@ -8,12 +8,16 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import type {
-  ESQLColumn,
-  ESQLCommand,
-  ESQLAstItem,
-  ESQLMessage,
-  ESQLFunction,
+import {
+  type ESQLColumn,
+  type ESQLCommand,
+  type ESQLAstItem,
+  type ESQLMessage,
+  type ESQLFunction,
+  isFunctionExpression,
+  isWhereExpression,
+  isFieldExpression,
+  Walker,
 } from '@kbn/esql-ast';
 import {
   getFunctionDefinition,
@@ -61,7 +65,13 @@ const statsValidator = (command: ESQLCommand) => {
   // until an agg function is detected
   // in the long run this might be integrated into the validation function
   const statsArg = command.args
-    .flatMap((arg) => (isAssignment(arg) ? arg.args[1] : arg))
+    .flatMap((arg) => {
+      if (isWhereExpression(arg) && isFunctionExpression(arg.args[0])) {
+        arg = arg.args[0] as ESQLFunction;
+      }
+
+      return isAssignment(arg) ? arg.args[1] : arg;
+    })
     .filter(isFunctionItem);
 
   if (statsArg.length) {
@@ -73,14 +83,31 @@ const statsValidator = (command: ESQLCommand) => {
     }
 
     function checkAggExistence(arg: ESQLFunction): boolean {
+      if (isWhereExpression(arg)) {
+        return checkAggExistence(arg.args[0] as ESQLFunction);
+      }
+
+      if (isFieldExpression(arg)) {
+        const agg = arg.args[1];
+        const firstFunction = Walker.match(agg, { type: 'function' });
+
+        if (!firstFunction) {
+          return false;
+        }
+
+        return checkAggExistence(firstFunction as ESQLFunction);
+      }
+
       // TODO the grouping function check may not
       // hold true for all future cases
       if (isAggFunction(arg) || isFunctionOperatorParam(arg)) {
         return true;
       }
+
       if (isOtherFunction(arg)) {
         return (arg as ESQLFunction).args.filter(isFunctionItem).some(checkAggExistence);
       }
+
       return false;
     }
     // first check: is there an agg function somewhere?
@@ -153,6 +180,7 @@ const statsValidator = (command: ESQLCommand) => {
       }
     }
   }
+
   return messages;
 };
 export const commandDefinitions: Array<CommandDefinition<any>> = [

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

@@ -112,10 +112,17 @@ export function collectVariables(
       addToVariables(oldArg, newArg, fields, variables);
     })
     .on('visitFunctionCallExpression', (ctx) => {
-      if (ctx.node.name === '=') {
-        addVariableFromAssignment(ctx.node, variables, fields);
+      const node = ctx.node;
+
+      if (node.subtype === 'binary-expression' && node.name === 'where') {
+        ctx.visitArgument(0, undefined);
+        return;
+      }
+
+      if (node.name === '=') {
+        addVariableFromAssignment(node, variables, fields);
       } else {
-        addVariableFromExpression(ctx.node, queryString, variables, fields);
+        addVariableFromExpression(node, queryString, variables, fields);
       }
     })
     .on('visitCommandOption', (ctx) => {

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/test_suites/validation.command.join.ts

@@ -13,7 +13,7 @@ export const validationJoinCommandTestSuite = (setup: helpers.Setup) => {
   describe('validation', () => {
     describe('command', () => {
       describe('<LEFT | RIGHT | LOOKUP> JOIN <index> [ AS <alias> ] ON <condition> [, <condition> [, ...]]', () => {
-        describe('... <index> [ AS <alias> ]', () => {
+        describe('... <index> ...', () => {
           test('validates the most basic query', async () => {
             const { expectErrors } = await setup();
 
@@ -45,6 +45,8 @@ export const validationJoinCommandTestSuite = (setup: helpers.Setup) => {
             await expectErrors('FROM index | LEFT JOIN join_index_alias_2 ON stringField', []);
           });
         });
+
+        test.todo('... AS <alias> ...');
       });
     });
   });

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/test_suites/validation.command.stats.ts

@@ -146,6 +146,12 @@ export const validationStatsCommandTestSuite = (setup: helpers.Setup) => {
               `Aggregate function's parameters must be an attribute, literal or a non-aggregation function; found [count(*)] of type [long]`,
             ]);
           });
+
+          test('allows WHERE clause', async () => {
+            const { expectErrors } = await setup();
+
+            await expectErrors('FROM a_index | STATS var0 = avg(doubleField) WHERE 123', []);
+          });
         });
 
         describe('... BY <grouping>', () => {

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

@@ -1225,13 +1225,9 @@ function validateCommand(
                 currentCommandIndex,
               })
             );
-          }
-
-          if (isSettingItem(arg)) {
+          } else if (isSettingItem(arg)) {
             messages.push(...validateSetting(arg, commandDef.modes[0], command, references));
-          }
-
-          if (isOptionItem(arg)) {
+          } else if (isOptionItem(arg)) {
             messages.push(
               ...validateOption(
                 arg,
@@ -1240,15 +1236,13 @@ function validateCommand(
                 references
               )
             );
-          }
-          if (isColumnItem(arg) || isIdentifier(arg)) {
+          } else if (isColumnItem(arg) || isIdentifier(arg)) {
             if (command.name === 'stats' || command.name === 'inlinestats') {
               messages.push(errors.unknownAggFunction(arg));
             } else {
               messages.push(...validateColumnForCommand(arg, command.name, references));
             }
-          }
-          if (isTimeIntervalItem(arg)) {
+          } else if (isTimeIntervalItem(arg)) {
             messages.push(
               getMessageFromId({
                 messageId: 'unsupportedTypeForCommand',
@@ -1260,8 +1254,7 @@ function validateCommand(
                 locations: arg.location,
               })
             );
-          }
-          if (isSourceItem(arg)) {
+          } else if (isSourceItem(arg)) {
             messages.push(...validateSource(arg, command.name, references));
           }
         }