[ES|QL] Improve STATS command summary extraction (elastic#199796)

## Summary Partially addresses elastic#191812 - Correctly extracts summary from of fields from the `BY` clause of `STATS` command. - The `.summarize()` command now returns `newFields` and `usedFields` properties. The `newFields` is a list of newly created fields by the `STATS` command. The `usedFields` is a list of all fields which were used by the `STATS` command. - Improves parameter node handling. ### Example Extract all "new" and "used" fields from all `STATS` commands: ```ts const query = EsqlQuery.fromSrc('FROM index | STATS a = max(b), agg(c) BY d'); const summary = mutate.commands.stats.summarize(query); console.log(summary.newFields); // [ 'a', '`agg(c)`' ] console.log(summary.usedFields); // [ 'b', 'c', 'd' ] ``` ### Checklist Delete any items that are not applicable to this PR. - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [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 ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#_add_your_labels)
CAWilson94 · Nov 18, 2024 · b6219cf · b6219cf
1 parent 2942100
commit b6219cf
Show file tree

Hide file tree

Showing 4 changed files with 184 additions and 56 deletions.
diff --git a/packages/kbn-esql-ast/src/ast/helpers.ts b/packages/kbn-esql-ast/src/ast/helpers.ts
@@ -14,6 +14,7 @@ import type {
   ESQLFunction,
   ESQLIntegerLiteral,
   ESQLLiteral,
+  ESQLParamLiteral,
   ESQLProperNode,
 } from '../types';
 import { BinaryExpressionGroup } from './constants';
@@ -48,6 +49,9 @@ export const isIntegerLiteral = (node: unknown): node is ESQLIntegerLiteral =>
 export const isDoubleLiteral = (node: unknown): node is ESQLIntegerLiteral =>
   isLiteral(node) && node.literalType === 'double';
 
+export const isParamLiteral = (node: unknown): node is ESQLParamLiteral =>
+  isLiteral(node) && node.literalType === 'param';
+
 export const isColumn = (node: unknown): node is ESQLColumn =>
   isProperNode(node) && node.type === 'column';
 

diff --git a/packages/kbn-esql-ast/src/mutate/README.md b/packages/kbn-esql-ast/src/mutate/README.md
@@ -60,3 +60,21 @@ console.log(src); // FROM index METADATA _lang, _id
     - `.remove()` &mdash; Remove a `LIMIT` command by index.
     - `.set()` &mdash; Set the limit value of a specific `LIMIT` command.
     - `.upsert()` &mdash; Insert a `LIMIT` command, or update the limit value if it already exists.
+  - `.stats`
+    - `.list()` &mdash; List all `STATS` commands.
+    - `.byIndex()` &mdash; Find a `STATS` command by index.
+    - `.summarize()` &mdash; Summarize all `STATS` commands.
+    - `.summarizeCommand()` &mdash; Summarize a specific `STATS` command.
+
+
+## Examples
+
+Extract all "new" and "used" fields from all `STATS` commands:
+
+```ts
+const query = EsqlQuery.fromSrc('FROM index | STATS a = max(b), agg(c) BY d');
+const summary = mutate.commands.stats.summarize(query);
+
+console.log(summary.newFields);     // [ 'a', '`agg(c)`' ]
+console.log(summary.usedFields);    // [ 'b', 'c', 'd' ]
+```
diff --git a/packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts
@@ -111,7 +111,7 @@ describe('commands.stats', () => {
                 name: 'bar',
               },
             ],
-            fields: ['bar'],
+            usedFields: new Set(['bar']),
           },
         },
       });
@@ -155,7 +155,7 @@ describe('commands.stats', () => {
                 name: 'baz',
               },
             ],
-            fields: ['foo', 'bar', 'baz'],
+            usedFields: new Set(['foo', 'bar', 'baz']),
           },
         },
       });
@@ -172,15 +172,15 @@ describe('commands.stats', () => {
         aggregates: {
           foo: {
             field: 'foo',
-            fields: ['f1', 'f2'],
+            usedFields: new Set(['f1', 'f2']),
           },
           'a.b': {
             field: 'a.b',
-            fields: ['f3'],
+            usedFields: new Set(['f3']),
           },
         },
       });
-      expect(summary.fields).toEqual(new Set(['f1', 'f2', 'f3']));
+      expect(summary.usedFields).toEqual(new Set(['f1', 'f2', 'f3']));
     });
 
     it('can get de-duplicated list of used fields', () => {
@@ -190,7 +190,7 @@ describe('commands.stats', () => {
       const command = commands.stats.byIndex(query.ast, 0)!;
       const summary = commands.stats.summarizeCommand(query, command);
 
-      expect(summary.fields).toEqual(new Set(['f1', 'f2']));
+      expect(summary.usedFields).toEqual(new Set(['f1', 'f2']));
     });
 
     describe('params', () => {
@@ -204,14 +204,14 @@ describe('commands.stats', () => {
         expect(summary).toMatchObject({
           aggregates: {
             foo: {
-              fields: ['f1.?aha', '?nested.?param'],
+              usedFields: new Set(['f1.?aha', '?nested.?param']),
             },
             'a.b': {
-              fields: ['f1'],
+              usedFields: new Set(['f1']),
             },
           },
         });
-        expect(summary.fields).toEqual(new Set(['f1.?aha', '?nested.?param', 'f1']));
+        expect(summary.usedFields).toEqual(new Set(['f1.?aha', '?nested.?param', 'f1']));
       });
 
       it('can use params as destination field names', () => {
@@ -224,11 +224,11 @@ describe('commands.stats', () => {
         expect(summary).toMatchObject({
           aggregates: {
             '?dest': {
-              fields: ['asdf'],
+              usedFields: new Set(['asdf']),
             },
           },
         });
-        expect(summary.fields).toEqual(new Set(['asdf']));
+        expect(summary.usedFields).toEqual(new Set(['asdf']));
       });
     });
 
@@ -243,7 +243,7 @@ describe('commands.stats', () => {
         expect(summary.aggregates).toEqual({
           '`max(1)`': expect.any(Object),
         });
-        expect(summary.fields).toEqual(new Set(['abc']));
+        expect(summary.usedFields).toEqual(new Set(['abc']));
       });
 
       it('returns all "grouping" fields', () => {
@@ -257,12 +257,45 @@ describe('commands.stats', () => {
           '`max(1)`': expect.any(Object),
         });
         expect(summary.grouping).toMatchObject({
-          a: { type: 'column' },
-          b: { type: 'column' },
-          c: { type: 'column' },
+          a: expect.any(Object),
+          b: expect.any(Object),
+          c: expect.any(Object),
         });
       });
 
+      it('returns grouping destination fields', () => {
+        const src = 'FROM index | STATS max(1) BY a, b, c';
+        const query = EsqlQuery.fromSrc(src);
+
+        const command = commands.stats.byIndex(query.ast, 0)!;
+        const summary = commands.stats.summarizeCommand(query, command);
+
+        expect(summary.aggregates).toEqual({
+          '`max(1)`': expect.any(Object),
+        });
+        expect(summary.grouping).toMatchObject({
+          a: expect.any(Object),
+          b: expect.any(Object),
+          c: expect.any(Object),
+        });
+        expect(summary.usedFields).toEqual(new Set(['a', 'b', 'c']));
+      });
+
+      it('returns grouping "used" fields', () => {
+        const src = 'FROM index | STATS max(1) BY a, b, c';
+        const query = EsqlQuery.fromSrc(src);
+
+        const command = commands.stats.byIndex(query.ast, 0)!;
+        const summary = commands.stats.summarizeCommand(query, command);
+
+        expect(summary.grouping).toMatchObject({
+          a: expect.any(Object),
+          b: expect.any(Object),
+          c: expect.any(Object),
+        });
+        expect(summary.usedFields).toEqual(new Set(['a', 'b', 'c']));
+      });
+
       it('can have params and quoted fields in grouping', () => {
         const src = 'FROM index | STATS max(1) BY `a😎`, ?123, a.?b.?0.`😎`';
         const query = EsqlQuery.fromSrc(src);
@@ -274,9 +307,9 @@ describe('commands.stats', () => {
           '`max(1)`': expect.any(Object),
         });
         expect(summary.grouping).toMatchObject({
-          '`a😎`': { type: 'column' },
-          // '?123': { type: 'column' },
-          'a.?b.?0.`😎`': { type: 'column' },
+          '`a😎`': expect.any(Object),
+          // '?123': expect.any(Object),
+          'a.?b.?0.`😎`': expect.any(Object),
         });
       });
     });
@@ -293,23 +326,67 @@ describe('commands.stats', () => {
           aggregates: {
             '`agg()`': {
               field: '`agg()`',
-              fields: [],
+              usedFields: new Set(),
             },
           },
-          fields: new Set([]),
+          usedFields: new Set([]),
         },
         {
           aggregates: {
             '`max(a, b, c)`': {
               field: '`max(a, b, c)`',
-              fields: ['a', 'b', 'c'],
+              usedFields: new Set(['a', 'b', 'c']),
             },
             '`max2(d.e)`': {
               field: '`max2(d.e)`',
-              fields: ['d.e'],
+              usedFields: new Set(['d.e']),
             },
           },
-          fields: new Set(['a', 'b', 'c', 'd.e']),
+          usedFields: new Set(['a', 'b', 'c', 'd.e']),
+        },
+      ]);
+    });
+
+    it('return used fields from BY clause', () => {
+      const src = 'FROM index | STATS agg(1) BY x, y = z, i = max(agg(1, 2, 3, ttt))';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.stats.summarize(query);
+
+      expect(summary).toMatchObject([
+        {
+          usedFields: new Set(['x', 'z', 'ttt']),
+        },
+      ]);
+    });
+
+    it('correctly returns used fields', () => {
+      const src =
+        'FROM index | LIMIT 1 | STATS agg(a, b), agg(c, a), d = agg(e) | LIMIT 2 | STATS max(a, b, c), max2(d.e) BY x, y = z, i = max(agg(1, 2, 3, ttt))';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.stats.summarize(query);
+
+      expect(summary).toMatchObject([
+        {
+          usedFields: new Set(['a', 'b', 'c', 'e']),
+        },
+        {
+          usedFields: new Set(['a', 'b', 'c', 'd.e', 'x', 'z', 'ttt']),
+        },
+      ]);
+    });
+
+    it('correctly returns new fields', () => {
+      const src =
+        'FROM index | LIMIT 1 | STATS agg(a, b), agg(c, a), d = agg(e) | LIMIT 2 | STATS max(a, b, c), max2(d.e) BY x, y = z, i = max(agg(1, 2, 3, ttt))';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.stats.summarize(query);
+
+      expect(summary).toMatchObject([
+        {
+          newFields: new Set(['`agg(a, b)`', '`agg(c, a)`', 'd']),
+        },
+        {
+          newFields: new Set(['`max(a, b, c)`', '`max2(d.e)`', 'x', 'y', 'i']),
         },
       ]);
     });