[Lens] optimize percentiles fetching

x-pack/plugins/lens/common/expressions/rename_columns/rename_columns.test.ts

@@ -26,14 +26,18 @@ describe('rename_columns', () => {
     };
 
     const idMap = {
-      a: {
-        id: 'b',
-        label: 'Austrailia',
-      },
-      b: {
-        id: 'c',
-        label: 'Boomerang',
-      },
+      a: [
+        {
+          id: 'b',
+          label: 'Austrailia',
+        },
+      ],
+      b: [
+        {
+          id: 'c',
+          label: 'Boomerang',
+        },
+      ],
     };
 
     const result = await renameColumns.fn(
@@ -99,7 +103,7 @@ describe('rename_columns', () => {
     };
 
     const idMap = {
-      b: { id: 'c', label: 'Catamaran' },
+      b: [{ id: 'c', label: 'Catamaran' }],
     };
 
     const result = await renameColumns.fn(
@@ -149,6 +153,67 @@ describe('rename_columns', () => {
     `);
   });
 
+  it('should map to multiple original columns', async () => {
+    const input: Datatable = {
+      type: 'datatable',
+      columns: [{ id: 'b', name: 'B', meta: { type: 'number' } }],
+      rows: [{ b: 2 }, { b: 4 }, { b: 6 }, { b: 8 }],
+    };
+
+    const idMap = {
+      b: [
+        { id: 'c', label: 'Catamaran' },
+        { id: 'd', label: 'Dinghy' },
+      ],
+    };
+
+    const result = await renameColumns.fn(
+      input,
+      { idMap: JSON.stringify(idMap) },
+      createMockExecutionContext()
+    );
+
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "columns": Array [
+          Object {
+            "id": "c",
+            "meta": Object {
+              "type": "number",
+            },
+            "name": "Catamaran",
+          },
+          Object {
+            "id": "d",
+            "meta": Object {
+              "type": "number",
+            },
+            "name": "Dinghy",
+          },
+        ],
+        "rows": Array [
+          Object {
+            "c": 2,
+            "d": 2,
+          },
+          Object {
+            "c": 4,
+            "d": 4,
+          },
+          Object {
+            "c": 6,
+            "d": 6,
+          },
+          Object {
+            "c": 8,
+            "d": 8,
+          },
+        ],
+        "type": "datatable",
+      }
+    `);
+  });
+
   it('should rename date histograms', async () => {
     const input: Datatable = {
       type: 'datatable',
@@ -165,7 +230,7 @@ describe('rename_columns', () => {
     };
 
     const idMap = {
-      b: { id: 'c', label: 'Apple', operationType: 'date_histogram', sourceField: 'banana' },
+      b: [{ id: 'c', label: 'Apple', operationType: 'date_histogram', sourceField: 'banana' }],
     };
 
     const result = await renameColumns.fn(

x-pack/plugins/lens/common/expressions/rename_columns/rename_columns.ts

@@ -6,10 +6,10 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import type { RenameColumnsExpressionFunction } from './types';
+import type { MapToOriginalColumnsExpressionFunction } from './types';
 
-export const renameColumns: RenameColumnsExpressionFunction = {
-  name: 'lens_rename_columns',
+export const renameColumns: MapToOriginalColumnsExpressionFunction = {
+  name: 'lens_map_to_original_columns',
   type: 'datatable',
   help: i18n.translate('xpack.lens.functions.renameColumns.help', {
     defaultMessage: 'A helper to rename the columns of a datatable',
@@ -26,7 +26,7 @@ export const renameColumns: RenameColumnsExpressionFunction = {
   inputTypes: ['datatable'],
   async fn(...args) {
     /** Build optimization: prevent adding extra code into initial bundle **/
-    const { renameColumnFn } = await import('./rename_columns_fn');
-    return renameColumnFn(...args);
+    const { mapToOriginalColumns } = await import('./rename_columns_fn');
+    return mapToOriginalColumns(...args);
   },
 };

x-pack/plugins/lens/common/expressions/rename_columns/rename_columns_fn.ts

@@ -6,7 +6,7 @@
  */
 
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
-import type { OriginalColumn, RenameColumnsExpressionFunction } from './types';
+import type { OriginalColumn, MapToOriginalColumnsExpressionFunction } from './types';
 
 function getColumnName(originalColumn: OriginalColumn, newColumn: DatatableColumn) {
   if (originalColumn?.operationType === 'date_histogram') {
@@ -21,42 +21,43 @@ function getColumnName(originalColumn: OriginalColumn, newColumn: DatatableColum
   return originalColumn.label;
 }
 
-export const renameColumnFn: RenameColumnsExpressionFunction['fn'] = (
+export const mapToOriginalColumns: MapToOriginalColumnsExpressionFunction['fn'] = (
   data,
   { idMap: encodedIdMap }
 ) => {
-  const idMap = JSON.parse(encodedIdMap) as Record<string, OriginalColumn>;
+  const idMap = JSON.parse(encodedIdMap) as Record<string, OriginalColumn[]>;
 
   return {
     ...data,
     rows: data.rows.map((row) => {
       const mappedRow: Record<string, unknown> = {};
-      Object.entries(idMap).forEach(([fromId, toId]) => {
-        mappedRow[toId.id] = row[fromId];
-      });
 
       Object.entries(row).forEach(([id, value]) => {
         if (id in idMap) {
-          mappedRow[idMap[id].id] = value;
+          idMap[id].forEach(({ id: originalId }) => {
+            mappedRow[originalId] = value;
+          });
         } else {
           mappedRow[id] = value;
         }
       });
 
       return mappedRow;
     }),
-    columns: data.columns.map((column) => {
-      const mappedItem = idMap[column.id];
-
-      if (!mappedItem) {
-        return column;
-      }
-
-      return {
-        ...column,
-        id: mappedItem.id,
-        name: getColumnName(mappedItem, column),
-      };
-    }),
+    columns: data.columns
+      .map((column) => {
+        const originalColumns = idMap[column.id];
+
+        if (!originalColumns) {
+          return column;
+        }
+
+        return originalColumns.map((originalColumn) => ({
+          ...column,
+          id: originalColumn.id,
+          name: getColumnName(originalColumn, column),
+        }));
+      })
+      .flat(),
   };
 };

x-pack/plugins/lens/common/expressions/rename_columns/types.ts

@@ -12,8 +12,8 @@ export type OriginalColumn = { id: string; label: string } & (
   | { operationType: string; sourceField: never }
 );
 
-export type RenameColumnsExpressionFunction = ExpressionFunctionDefinition<
-  'lens_rename_columns',
+export type MapToOriginalColumnsExpressionFunction = ExpressionFunctionDefinition<
+  'lens_map_to_original_columns',
   Datatable,
   {
     idMap: string;

x-pack/plugins/lens/public/indexpattern_datasource/indexpattern.test.ts

@@ -491,10 +491,10 @@ describe('IndexPattern Data Source', () => {
             Object {
               "arguments": Object {
                 "idMap": Array [
-                  "{\\"col-0-0\\":{\\"label\\":\\"Count of records\\",\\"dataType\\":\\"number\\",\\"isBucketed\\":false,\\"sourceField\\":\\"___records___\\",\\"operationType\\":\\"count\\",\\"id\\":\\"col1\\"},\\"col-1-1\\":{\\"label\\":\\"Date\\",\\"dataType\\":\\"date\\",\\"isBucketed\\":true,\\"operationType\\":\\"date_histogram\\",\\"sourceField\\":\\"timestamp\\",\\"params\\":{\\"interval\\":\\"1d\\"},\\"id\\":\\"col2\\"}}",
+                  "{\\"col-0-0\\":[{\\"label\\":\\"Count of records\\",\\"dataType\\":\\"number\\",\\"isBucketed\\":false,\\"sourceField\\":\\"___records___\\",\\"operationType\\":\\"count\\",\\"id\\":\\"col1\\"}],\\"col-1-1\\":[{\\"label\\":\\"Date\\",\\"dataType\\":\\"date\\",\\"isBucketed\\":true,\\"operationType\\":\\"date_histogram\\",\\"sourceField\\":\\"timestamp\\",\\"params\\":{\\"interval\\":\\"1d\\"},\\"id\\":\\"col2\\"}]}",
                 ],
               },
-              "function": "lens_rename_columns",
+              "function": "lens_map_to_original_columns",
               "type": "function",
             },
           ],
@@ -905,9 +905,9 @@ describe('IndexPattern Data Source', () => {
       const ast = indexPatternDatasource.toExpression(state, 'first') as Ast;
       expect(ast.chain[1].arguments.metricsAtAllLevels).toEqual([false]);
       expect(JSON.parse(ast.chain[2].arguments.idMap[0] as string)).toEqual({
-        'col-0-0': expect.objectContaining({ id: 'bucket1' }),
-        'col-1-1': expect.objectContaining({ id: 'bucket2' }),
-        'col-2-2': expect.objectContaining({ id: 'metric' }),
+        'col-0-0': [expect.objectContaining({ id: 'bucket1' })],
+        'col-1-1': [expect.objectContaining({ id: 'bucket2' })],
+        'col-2-2': [expect.objectContaining({ id: 'metric' })],
       });
     });
 
@@ -1026,10 +1026,13 @@ describe('IndexPattern Data Source', () => {
         const state = enrichBaseState(queryBaseState);
 
         const ast = indexPatternDatasource.toExpression(state, 'first') as Ast;
+
         expect(JSON.parse(ast.chain[2].arguments.idMap[0] as string)).toEqual({
-          'col-0-0': expect.objectContaining({
-            id: 'col1',
-          }),
+          'col-0-0': [
+            expect.objectContaining({
+              id: 'col1',
+            }),
+          ],
         });
       });
 

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/index.ts

@@ -12,7 +12,10 @@ import {
   CoreStart,
 } from '@kbn/core/public';
 import { IStorageWrapper } from '@kbn/kibana-utils-plugin/public';
-import { ExpressionAstFunction } from '@kbn/expressions-plugin/public';
+import {
+  ExpressionAstExpressionBuilder,
+  ExpressionAstFunction,
+} from '@kbn/expressions-plugin/public';
 import { DataPublicPluginStart } from '@kbn/data-plugin/public';
 import { UnifiedSearchPublicPluginStart } from '@kbn/unified-search-plugin/public';
 import { DataViewsPublicPluginStart } from '@kbn/data-views-plugin/public';
@@ -55,6 +58,7 @@ import { IndexPattern, IndexPatternField, IndexPatternLayer } from '../../types'
 import { DateRange, LayerType } from '../../../../common';
 import { rangeOperation } from './ranges';
 import { IndexPatternDimensionEditorProps, OperationSupportMatrix } from '../../dimension_panel';
+import type { OriginalColumn } from '../../to_expression';
 
 export type {
   IncompleteColumn,
@@ -378,6 +382,17 @@ interface BaseOperationDefinitionProps<C extends BaseIndexPatternColumn, P = {}>
    * Title for the help component
    */
   helpComponentTitle?: string;
+  /**
+   * Optimizes EsAggs expression. Invoked only once per operation type. May mutate arguments.
+   */
+  optimizeEsAggs?: (
+    aggs: ExpressionAstExpressionBuilder[],
+    esAggsIdMap: Record<string, OriginalColumn[]>,
+    aggExpressionToEsAggsIdMap: Map<ExpressionAstExpressionBuilder, string>
+  ) => {
+    aggs: ExpressionAstExpressionBuilder[];
+    esAggsIdMap: Record<string, OriginalColumn[]>;
+  };
 }
 
 interface BaseBuildColumnArgs {