[ES|QL] Implement OrderExpression for SORT command arguments

packages/kbn-esql-ast/src/__tests__/ast_parser.sort.test.ts

@@ -11,34 +11,107 @@ import { getAstAndSyntaxErrors as parse } from '../ast_parser';
 
 describe('SORT', () => {
   describe('correctly formatted', () => {
-    // Un-skip one https://github.com/elastic/kibana/issues/189491 fixed.
-    it.skip('example from documentation', () => {
+    it('sorting order without modifiers', () => {
+      const text = `FROM employees | SORT height`;
+      const { ast, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'sort',
+          args: [
+            {
+              type: 'column',
+              name: 'height',
+            },
+          ],
+        },
+      ]);
+    });
+
+    it('sort expression is a function call', () => {
+      const text = `from a_index | sort values(textField)`;
+      const { ast, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'sort',
+          args: [
+            {
+              type: 'function',
+              name: 'values',
+            },
+          ],
+        },
+      ]);
+    });
+
+    it('with order modifier "DESC"', () => {
       const text = `
         FROM employees
-        | KEEP first_name, last_name, height
         | SORT height DESC
         `;
       const { ast, errors } = parse(text);
 
       expect(errors.length).toBe(0);
       expect(ast).toMatchObject([
         {},
+        {
+          type: 'command',
+          name: 'sort',
+          args: [
+            {
+              type: 'order',
+              order: 'DESC',
+              nulls: '',
+              args: [
+                {
+                  type: 'column',
+                  name: 'height',
+                },
+              ],
+            },
+          ],
+        },
+      ]);
+    });
+
+    it('with nulls modifier "NULLS LAST"', () => {
+      const text = `
+        FROM employees
+        | SORT height NULLS LAST
+        `;
+      const { ast, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast).toMatchObject([
         {},
         {
           type: 'command',
           name: 'sort',
           args: [
             {
-              type: 'column',
-              name: 'height',
+              type: 'order',
+              order: '',
+              nulls: 'NULLS LAST',
+              args: [
+                {
+                  type: 'column',
+                  name: 'height',
+                },
+              ],
             },
           ],
         },
       ]);
     });
 
-    // Un-skip once https://github.com/elastic/kibana/issues/189491 fixed.
-    it.skip('can parse various sorting columns with options', () => {
+    it('can parse various sorting columns with options', () => {
       const text =
         'FROM a | SORT a, b ASC, c DESC, d NULLS FIRST, e NULLS LAST, f ASC NULLS FIRST, g DESC NULLS LAST';
       const { ast, errors } = parse(text);
@@ -55,28 +128,58 @@ describe('SORT', () => {
               name: 'a',
             },
             {
-              type: 'column',
-              name: 'b',
+              order: 'ASC',
+              nulls: '',
+              args: [
+                {
+                  name: 'b',
+                },
+              ],
             },
             {
-              type: 'column',
-              name: 'c',
+              order: 'DESC',
+              nulls: '',
+              args: [
+                {
+                  name: 'c',
+                },
+              ],
             },
             {
-              type: 'column',
-              name: 'd',
+              order: '',
+              nulls: 'NULLS FIRST',
+              args: [
+                {
+                  name: 'd',
+                },
+              ],
             },
             {
-              type: 'column',
-              name: 'e',
+              order: '',
+              nulls: 'NULLS LAST',
+              args: [
+                {
+                  name: 'e',
+                },
+              ],
             },
             {
-              type: 'column',
-              name: 'f',
+              order: 'ASC',
+              nulls: 'NULLS FIRST',
+              args: [
+                {
+                  name: 'f',
+                },
+              ],
             },
             {
-              type: 'column',
-              name: 'g',
+              order: 'DESC',
+              nulls: 'NULLS LAST',
+              args: [
+                {
+                  name: 'g',
+                },
+              ],
             },
           ],
         },

packages/kbn-esql-ast/src/ast_factory.ts

@@ -53,7 +53,7 @@ import {
   visitDissect,
   visitGrok,
   collectBooleanExpression,
-  visitOrderExpression,
+  visitOrderExpressions,
   getPolicyName,
   getMatchField,
   getEnrichClauses,
@@ -238,7 +238,7 @@ export class AstListener implements ESQLParserListener {
   exitSortCommand(ctx: SortCommandContext) {
     const command = createCommand('sort', ctx);
     this.ast.push(command);
-    command.args.push(...visitOrderExpression(ctx.orderExpression_list()));
+    command.args.push(...visitOrderExpressions(ctx.orderExpression_list()));
   }
 
   /**

packages/kbn-esql-ast/src/ast_helpers.ts

@@ -42,6 +42,7 @@ import type {
   ESQLNumericLiteralType,
   FunctionSubtype,
   ESQLNumericLiteral,
+  ESQLOrderExpression,
 } from './types';
 import { parseIdentifier } from './parser/helpers';
 
@@ -222,6 +223,26 @@ export function createFunction<Subtype extends FunctionSubtype>(
   return node;
 }
 
+export const createOrderExpression = (
+  ctx: ParserRuleContext,
+  arg: ESQLAstItem,
+  order: ESQLOrderExpression['order'],
+  nulls: ESQLOrderExpression['nulls']
+) => {
+  const node: ESQLOrderExpression = {
+    type: 'order',
+    name: '',
+    order,
+    nulls,
+    args: [arg],
+    text: ctx.getText(),
+    location: getPosition(ctx.start, ctx.stop),
+    incomplete: Boolean(ctx.exception),
+  };
+
+  return node;
+};
+
 function walkFunctionStructure(
   args: ESQLAstItem[],
   initialLocation: ESQLLocation,

packages/kbn-esql-ast/src/ast_walker.ts

@@ -84,6 +84,7 @@ import {
   textExistsAndIsValid,
   createInlineCast,
   createUnknownItem,
+  createOrderExpression,
 } from './ast_helpers';
 import { getPosition } from './ast_position_utils';
 import {
@@ -97,6 +98,7 @@ import {
   ESQLUnnamedParamLiteral,
   ESQLPositionalParamLiteral,
   ESQLNamedParamLiteral,
+  ESQLOrderExpression,
 } from './types';
 
 export function collectAllSourceIdentifiers(ctx: FromCommandContext): ESQLAstItem[] {
@@ -608,34 +610,43 @@ export function visitByOption(
   return [option];
 }
 
-export function visitOrderExpression(ctx: OrderExpressionContext[]) {
-  const ast: ESQLAstItem[] = [];
-  for (const orderCtx of ctx) {
-    const expression = collectBooleanExpression(orderCtx.booleanExpression());
-    if (orderCtx._ordering) {
-      const terminalNode =
-        orderCtx.getToken(esql_parser.ASC, 0) || orderCtx.getToken(esql_parser.DESC, 0);
-      const literal = createLiteral('string', terminalNode);
-      if (literal) {
-        expression.push(literal);
-      }
-    }
-    if (orderCtx.NULLS()) {
-      expression.push(createLiteral('string', orderCtx.NULLS()!)!);
-      if (orderCtx._nullOrdering && orderCtx._nullOrdering.text !== '<first missing>') {
-        const innerTerminalNode =
-          orderCtx.getToken(esql_parser.FIRST, 0) || orderCtx.getToken(esql_parser.LAST, 0);
-        const literal = createLiteral('string', innerTerminalNode);
-        if (literal) {
-          expression.push(literal);
-        }
-      }
-    }
+const visitOrderExpression = (ctx: OrderExpressionContext): ESQLOrderExpression | ESQLAstItem => {
+  const arg = collectBooleanExpression(ctx.booleanExpression())[0];
 
-    if (expression.length) {
-      ast.push(...expression);
-    }
+  let order: ESQLOrderExpression['order'] = '';
+  let nulls: ESQLOrderExpression['nulls'] = '';
+
+  const ordering = ctx._ordering?.text?.toUpperCase();
+
+  if (ordering) order = ordering as ESQLOrderExpression['order'];
+
+  const nullOrdering = ctx._nullOrdering?.text?.toUpperCase();
+
+  switch (nullOrdering) {
+    case 'LAST':
+      nulls = 'NULLS LAST';
+      break;
+    case 'FIRST':
+      nulls = 'NULLS FIRST';
+      break;
   }
+
+  if (!order && !nulls) {
+    return arg;
+  }
+
+  return createOrderExpression(ctx, arg, order, nulls);
+};
+
+export function visitOrderExpressions(
+  ctx: OrderExpressionContext[]
+): Array<ESQLOrderExpression | ESQLAstItem> {
+  const ast: Array<ESQLOrderExpression | ESQLAstItem> = [];
+
+  for (const orderCtx of ctx) {
+    ast.push(visitOrderExpression(orderCtx));
+  }
+
   return ast;
 }
 

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

@@ -50,18 +50,22 @@ describe('single line query', () => {
         expect(text).toBe('FROM a | SORT b');
       });
 
-      /** @todo Enable once order expressions are supported.  */
-      test.skip('order expression with ASC modifier', () => {
+      test('order expression with ASC modifier', () => {
         const { text } = reprint('FROM a | SORT b ASC');
 
         expect(text).toBe('FROM a | SORT b ASC');
       });
 
-      /** @todo Enable once order expressions are supported.  */
-      test.skip('order expression with ASC and NULLS FIRST modifier', () => {
-        const { text } = reprint('FROM a | SORT b ASC NULLS FIRST');
+      test('order expression with NULLS LAST modifier', () => {
+        const { text } = reprint('FROM a | SORT b NULLS LAST');
 
-        expect(text).toBe('FROM a | SORT b ASC NULLS FIRST');
+        expect(text).toBe('FROM a | SORT b NULLS LAST');
+      });
+
+      test('order expression with DESC and NULLS FIRST modifier', () => {
+        const { text } = reprint('FROM a | SORT b DESC NULLS FIRST');
+
+        expect(text).toBe('FROM a | SORT b DESC NULLS FIRST');
       });
     });