[Lens] Rename columns

x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.test.tsx

@@ -225,7 +225,7 @@ describe('IndexPattern Data Source', () => {
       };
       const state = await indexPatternDatasource.initialize(queryPersistedState);
       expect(indexPatternDatasource.toExpression(state)).toMatchInlineSnapshot(
-        `"esdocs index=\\"my-fake-index-pattern\\" fields=\\"source, bytes\\" sort=\\"source, DESC\\""`
+        `"esdocs index=\\"my-fake-index-pattern\\" fields=\\"source, bytes\\" sort=\\"source, DESC\\" | lens_rename_columns idMap='{\\"source\\":\\"col1\\",\\"bytes\\":\\"col2\\"}'"`
       );
     });
 
@@ -262,7 +262,7 @@ describe('IndexPattern Data Source', () => {
       index=\\"1\\"
       metricsAtAllLevels=\\"false\\"
       partialRows=\\"false\\"
-      aggConfigs='[{\\"id\\":\\"col1\\",\\"enabled\\":true,\\"type\\":\\"count\\",\\"schema\\":\\"metric\\",\\"params\\":{}},{\\"id\\":\\"col2\\",\\"enabled\\":true,\\"type\\":\\"date_histogram\\",\\"schema\\":\\"segment\\",\\"params\\":{\\"field\\":\\"timestamp\\",\\"timeRange\\":{\\"from\\":\\"now-1d\\",\\"to\\":\\"now\\"},\\"useNormalizedEsInterval\\":true,\\"interval\\":\\"1h\\",\\"drop_partials\\":false,\\"min_doc_count\\":1,\\"extended_bounds\\":{}}}]'"
+      aggConfigs='[{\\"id\\":\\"col1\\",\\"enabled\\":true,\\"type\\":\\"count\\",\\"schema\\":\\"metric\\",\\"params\\":{}},{\\"id\\":\\"col2\\",\\"enabled\\":true,\\"type\\":\\"date_histogram\\",\\"schema\\":\\"segment\\",\\"params\\":{\\"field\\":\\"timestamp\\",\\"timeRange\\":{\\"from\\":\\"now-1d\\",\\"to\\":\\"now\\"},\\"useNormalizedEsInterval\\":true,\\"interval\\":\\"1h\\",\\"drop_partials\\":false,\\"min_doc_count\\":1,\\"extended_bounds\\":{}}}]' | lens_rename_columns idMap='{\\"col-0-col1\\":\\"col1\\",\\"col-1-col2\\":\\"col2\\"}'"
 `);
     });
   });

x-pack/plugins/lens/public/indexpattern_plugin/plugin.tsx

@@ -4,14 +4,39 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+import { Registry } from '@kbn/interpreter/target/common';
+import { CoreSetup } from 'src/core/public';
 import chrome from 'ui/chrome';
 import { toastNotifications } from 'ui/notify';
 import { getIndexPatternDatasource } from './indexpattern';
 
+import {
+  functionsRegistry,
+  // @ts-ignore untyped dependency
+} from '../../../../../src/legacy/core_plugins/interpreter/public/registries';
+import { ExpressionFunction } from '../../../../../src/legacy/core_plugins/interpreter/public';
+import { renameColumns } from './rename_columns';
+
+// TODO these are intermediary types because interpreter is not typed yet
+// They can get replaced by references to the real interfaces as soon as they
+// are available
+
+export interface IndexPatternDatasourcePluginPlugins {
+  interpreter: InterpreterSetup;
+}
+
+export interface InterpreterSetup {
+  functionsRegistry: Registry<
+    ExpressionFunction<string, unknown, unknown, unknown>,
+    ExpressionFunction<string, unknown, unknown, unknown>
+  >;
+}
+
 class IndexPatternDatasourcePlugin {
   constructor() {}
 
-  setup() {
+  setup(_core: CoreSetup | null, { interpreter }: IndexPatternDatasourcePluginPlugins) {
+    interpreter.functionsRegistry.register(() => renameColumns);
     return getIndexPatternDatasource(chrome, toastNotifications);
   }
 
@@ -20,5 +45,10 @@ class IndexPatternDatasourcePlugin {
 
 const plugin = new IndexPatternDatasourcePlugin();
 
-export const indexPatternDatasourceSetup = () => plugin.setup();
+export const indexPatternDatasourceSetup = () =>
+  plugin.setup(null, {
+    interpreter: {
+      functionsRegistry,
+    },
+  });
 export const indexPatternDatasourceStop = () => plugin.stop();

x-pack/plugins/lens/public/indexpattern_plugin/rename_columns.test.ts

@@ -0,0 +1,103 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { renameColumns } from './rename_columns';
+import { KibanaDatatable } from 'src/legacy/core_plugins/interpreter/types';
+
+describe('rename_columns', () => {
+  it('should rename columns of a given datatable', () => {
+    const input: KibanaDatatable = {
+      type: 'kibana_datatable',
+      columns: [{ id: 'a', name: 'A' }, { id: 'b', name: 'B' }],
+      rows: [{ a: 1, b: 2 }, { a: 3, b: 4 }, { a: 5, b: 6 }, { a: 7, b: 8 }],
+    };
+
+    const idMap = {
+      a: 'b',
+      b: 'c',
+    };
+
+    expect(renameColumns.fn(input, { idMap: JSON.stringify(idMap) }, {})).toMatchInlineSnapshot(`
+Object {
+  "columns": Array [
+    Object {
+      "id": "b",
+      "name": "A",
+    },
+    Object {
+      "id": "c",
+      "name": "B",
+    },
+  ],
+  "rows": Array [
+    Object {
+      "b": 1,
+      "c": 2,
+    },
+    Object {
+      "b": 3,
+      "c": 4,
+    },
+    Object {
+      "b": 5,
+      "c": 6,
+    },
+    Object {
+      "b": 7,
+      "c": 8,
+    },
+  ],
+  "type": "kibana_datatable",
+}
+`);
+  });
+
+  it('should keep columns which are not mapped', () => {
+    const input: KibanaDatatable = {
+      type: 'kibana_datatable',
+      columns: [{ id: 'a', name: 'A' }, { id: 'b', name: 'B' }],
+      rows: [{ a: 1, b: 2 }, { a: 3, b: 4 }, { a: 5, b: 6 }, { a: 7, b: 8 }],
+    };
+
+    const idMap = {
+      b: 'c',
+    };
+
+    expect(renameColumns.fn(input, { idMap: JSON.stringify(idMap) }, {})).toMatchInlineSnapshot(`
+Object {
+  "columns": Array [
+    Object {
+      "id": "a",
+      "name": "A",
+    },
+    Object {
+      "id": "c",
+      "name": "B",
+    },
+  ],
+  "rows": Array [
+    Object {
+      "a": 1,
+      "c": 2,
+    },
+    Object {
+      "a": 3,
+      "c": 4,
+    },
+    Object {
+      "a": 5,
+      "c": 6,
+    },
+    Object {
+      "a": 7,
+      "c": 8,
+    },
+  ],
+  "type": "kibana_datatable",
+}
+`);
+  });
+});

x-pack/plugins/lens/public/indexpattern_plugin/rename_columns.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;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { i18n } from '@kbn/i18n';
+import { ExpressionFunction } from 'src/legacy/core_plugins/interpreter/types';
+import { KibanaDatatable } from '../types';
+
+interface RemapArgs {
+  idMap: string;
+}
+
+export const renameColumns: ExpressionFunction<
+  'lens_rename_columns',
+  KibanaDatatable,
+  RemapArgs,
+  KibanaDatatable
+> = {
+  name: 'lens_rename_columns',
+  type: 'kibana_datatable',
+  help: i18n.translate('lens.functions.renameColumns.help', {
+    defaultMessage: 'A helper to rename the columns of a datatable',
+  }),
+  args: {
+    idMap: {
+      types: ['string'],
+      help: i18n.translate('lens.functions.renameColumns.idMap.help', {
+        defaultMessage:
+          'A JSON encoded object in which keys are the old column ids and values are the corresponding new ones. All other columns ids are kept.',
+      }),
+    },
+  },
+  context: {
+    types: ['kibana_datatable'],
+  },
+  fn(data: KibanaDatatable, { idMap: encodedIdMap }: RemapArgs) {
+    const idMap = JSON.parse(encodedIdMap) as Record<string, string>;
+    return {
+      type: 'kibana_datatable',
+      rows: data.rows.map(row => {
+        const mappedRow: Record<string, unknown> = {};
+        Object.entries(idMap).forEach(([fromId, toId]) => {
+          mappedRow[toId] = row[fromId];
+        });
+
+        Object.entries(row).forEach(([id, value]) => {
+          if (id in idMap) {
+            mappedRow[idMap[id]] = value;
+          } else {
+            mappedRow[id] = value;
+          }
+        });
+
+        return mappedRow;
+      }),
+      columns: data.columns.map(column => ({
+        ...column,
+        id: idMap[column.id] ? idMap[column.id] : column.id,
+      })),
+    };
+  },
+};

x-pack/plugins/lens/public/indexpattern_plugin/to_expression.ts

@@ -19,9 +19,16 @@ export function toExpression(state: IndexPatternPrivateState) {
   const indexName = state.indexPatterns[state.currentIndexPatternId].title;
 
   if (sortedColumns.every(({ operationType }) => operationType === 'value')) {
+    const idMap = fieldNames.reduce(
+      (currentIdMap, fieldName, index) => ({
+        ...currentIdMap,
+        [fieldName]: state.columnOrder[index],
+      }),
+      {} as Record<string, string>
+    );
     return `esdocs index="${indexName}" fields="${fieldNames.join(', ')}" sort="${
       fieldNames[0]
-    }, DESC"`;
+    }, DESC" | lens_rename_columns idMap='${JSON.stringify(idMap)}'`;
   } else if (sortedColumns.length) {
     const firstMetric = sortedColumns.findIndex(({ isBucketed }) => !isBucketed);
     const aggs = sortedColumns.map((col, index) => {
@@ -83,12 +90,20 @@ export function toExpression(state: IndexPatternPrivateState) {
       }
     });
 
+    const idMap = state.columnOrder.reduce(
+      (currentIdMap, columnId, index) => ({
+        ...currentIdMap,
+        [`col-${index}-${columnId}`]: columnId,
+      }),
+      {} as Record<string, string>
+    );
+
     return `esaggs
       index="${state.currentIndexPatternId}"
       metricsAtAllLevels="false"
       partialRows="false"
-      aggConfigs='${JSON.stringify(aggs)}'`;
+      aggConfigs='${JSON.stringify(aggs)}' | lens_rename_columns idMap='${JSON.stringify(idMap)}'`;
   }
 
-  return '';
+  return null;
 }