[Lens] optimize percentiles fetching (#131875)

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

@@ -8,6 +8,6 @@
 export * from './counter_rate';
 export * from './collapse';
 export * from './format_column';
-export * from './rename_columns';
+export * from './map_to_columns';
 export * from './time_scale';
 export * from './datatable';

x-pack/plugins/lens/common/expressions/rename_columns/index.ts → x-pack/plugins/lens/common/expressions/map_to_columns/index.ts

@@ -5,4 +5,4 @@
  * 2.0.
  */
 
-export { renameColumns } from './rename_columns';
+export { mapToColumns } from './map_to_columns';

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

@@ -5,11 +5,11 @@
  * 2.0.
  */
 
-import { renameColumns } from './rename_columns';
+import { mapToColumns } from './map_to_columns';
 import { Datatable } from '@kbn/expressions-plugin/common';
 import { createMockExecutionContext } from '@kbn/expressions-plugin/common/mocks';
 
-describe('rename_columns', () => {
+describe('map_to_columns', () => {
   it('should rename columns of a given datatable', async () => {
     const input: Datatable = {
       type: 'datatable',
@@ -26,17 +26,21 @@ 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(
+    const result = await mapToColumns.fn(
       input,
       { idMap: JSON.stringify(idMap) },
       createMockExecutionContext()
@@ -99,10 +103,10 @@ describe('rename_columns', () => {
     };
 
     const idMap = {
-      b: { id: 'c', label: 'Catamaran' },
+      b: [{ id: 'c', label: 'Catamaran' }],
     };
 
-    const result = await renameColumns.fn(
+    const result = await mapToColumns.fn(
       input,
       { idMap: JSON.stringify(idMap) },
       createMockExecutionContext()
@@ -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 mapToColumns.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,10 +230,10 @@ 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(
+    const result = await mapToColumns.fn(
       input,
       { idMap: JSON.stringify(idMap) },
       createMockExecutionContext()

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

@@ -0,0 +1,32 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { i18n } from '@kbn/i18n';
+import type { MapToColumnsExpressionFunction } from './types';
+
+export const mapToColumns: MapToColumnsExpressionFunction = {
+  name: 'lens_map_to_columns',
+  type: 'datatable',
+  help: i18n.translate('xpack.lens.functions.mapToColumns.help', {
+    defaultMessage: 'A helper to transform a datatable to match Lens column definitions',
+  }),
+  args: {
+    idMap: {
+      types: ['string'],
+      help: i18n.translate('xpack.lens.functions.mapToColumns.idMap.help', {
+        defaultMessage:
+          'A JSON encoded object in which keys are the datatable column ids and values are the Lens column definitions. Any datatable columns not mentioned within the ID map will be kept unmapped.',
+      }),
+    },
+  },
+  inputTypes: ['datatable'],
+  async fn(...args) {
+    /** Build optimization: prevent adding extra code into initial bundle **/
+    const { mapToOriginalColumns } = await import('./map_to_columns_fn');
+    return mapToOriginalColumns(...args);
+  },
+};

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

@@ -6,7 +6,7 @@
  */
 
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
-import type { OriginalColumn, RenameColumnsExpressionFunction } from './types';
+import type { OriginalColumn, MapToColumnsExpressionFunction } 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: MapToColumnsExpressionFunction['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 → x-pack/plugins/lens/common/expressions/map_to_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 MapToColumnsExpressionFunction = ExpressionFunctionDefinition<
+  'lens_map_to_columns',
   Datatable,
   {
     idMap: string;

x-pack/plugins/lens/public/expressions.ts

@@ -8,7 +8,7 @@
 import type { ExpressionsSetup } from '@kbn/expressions-plugin/public';
 import { getDatatable } from '../common/expressions/datatable/datatable';
 import { datatableColumn } from '../common/expressions/datatable/datatable_column';
-import { renameColumns } from '../common/expressions/rename_columns/rename_columns';
+import { mapToColumns } from '../common/expressions/map_to_columns/map_to_columns';
 import { formatColumn } from '../common/expressions/format_column';
 import { counterRate } from '../common/expressions/counter_rate';
 import { getTimeScale } from '../common/expressions/time_scale/time_scale';
@@ -24,7 +24,7 @@ export const setupExpressions = (
     collapse,
     counterRate,
     formatColumn,
-    renameColumns,
+    mapToColumns,
     datatableColumn,
     getDatatable(formatFactory),
     getTimeScale(getDatatableUtilities, getTimeZone),