[Lens][Table] Fix csv export column sort order (#236673)

## Summary

When a table from a dashboard using the **Download CSV** action, the
column sort order visible in the table on the dashboard are preserved in
the csv output.


https://github.com/user-attachments/assets/7454a511-150b-45a2-86c5-1cd61bc0191a

Fixes #236550

## Details

The `datatable_fn` is used to set the table of data to the `adapters`,
which is the data used in the csv export action. This data comes in out
of order with additional `Part of X` columns for formula columns.
However, we do not sort these in before passing to the adapters.
Normally, with formulas, this is not a problem as the columns are in the
correct order but the
[`tabify`](https://github.com/elastic/kibana/blob/8d4b0956586284c35db97de2baea49b58476ee2a/src/platform/plugins/shared/data/common/search/tabify/tabify.ts#L24)
logic mixes up the column order for formulas.

The fix is to simply sort the columns in the table before passing it to
the `adapters.`

### Checklist

- [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
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)
- [x] Review the [backport
guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing)
and apply applicable `backport:*` labels.

## Release Note

Fixes a bug in the Lens table in which exporting the table from a
dashboard, which containg formula columns, can result in a different
column order than shown on the dashboard.

Co-authored-by: Marco Liberati <[email protected]>

Author: nickofthyme

x-pack/platform/plugins/shared/lens/common/expressions/impl/datatable/datatable_fn.test.ts

@@ -0,0 +1,122 @@
+/*
+ * 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 { fieldFormatsServiceMock } from '@kbn/field-formats-plugin/public/mocks';
+import type {
+  Datatable,
+  DefaultInspectorAdapters,
+  ExecutionContext,
+} from '@kbn/expressions-plugin/common';
+import type { DatatableArgs } from '../..';
+import { datatableFn } from './datatable_fn';
+import { shuffle } from 'lodash';
+
+const context = {
+  variables: { embeddableTitle: 'title' },
+} as unknown as ExecutionContext<DefaultInspectorAdapters>;
+
+const mockFormatFactory = fieldFormatsServiceMock.createStartContract().deserialize;
+
+describe('datatableFn', () => {
+  function buildTable(): Datatable {
+    return {
+      type: 'datatable',
+      columns: [
+        { id: 'bucket1', name: 'bucket1', meta: { type: 'string' } },
+        { id: 'bucket2', name: 'bucket2', meta: { type: 'string' } },
+        { id: 'bucket3', name: 'bucket3', meta: { type: 'string' } },
+        { id: 'metric1', name: 'metric1', meta: { type: 'number' } },
+        { id: 'metric2', name: 'metric2', meta: { type: 'number' } },
+      ],
+      rows: [
+        { bucket1: 'A', bucket2: 'D', bucket3: 'X', metric1: 1, metric2: 2 },
+        { bucket1: 'A', bucket2: 'D', bucket3: 'Y', metric1: 3, metric2: 4 },
+        { bucket1: 'A', bucket2: 'D', bucket3: 'Z', metric1: 5, metric2: 6 },
+        { bucket1: 'A', bucket2: 'E', bucket3: 'X', metric1: 7, metric2: 8 },
+        { bucket1: 'A', bucket2: 'E', bucket3: 'Y', metric1: 9, metric2: 10 },
+        { bucket1: 'A', bucket2: 'E', bucket3: 'Z', metric1: 11, metric2: 12 },
+        { bucket1: 'A', bucket2: 'F', bucket3: 'X', metric1: 13, metric2: 14 },
+        { bucket1: 'A', bucket2: 'F', bucket3: 'Y', metric1: 15, metric2: 16 },
+        { bucket1: 'A', bucket2: 'F', bucket3: 'Z', metric1: 17, metric2: 18 },
+        { bucket1: 'B', bucket2: 'D', bucket3: 'X', metric1: 19, metric2: 20 },
+        { bucket1: 'B', bucket2: 'D', bucket3: 'Y', metric1: 21, metric2: 22 },
+        { bucket1: 'B', bucket2: 'D', bucket3: 'Z', metric1: 23, metric2: 24 },
+        { bucket1: 'B', bucket2: 'E', bucket3: 'X', metric1: 25, metric2: 26 },
+        { bucket1: 'B', bucket2: 'E', bucket3: 'Y', metric1: 27, metric2: 28 },
+        { bucket1: 'B', bucket2: 'E', bucket3: 'Z', metric1: 29, metric2: 30 },
+        { bucket1: 'B', bucket2: 'F', bucket3: 'X', metric1: 31, metric2: 32 },
+        { bucket1: 'B', bucket2: 'F', bucket3: 'Y', metric1: 33, metric2: 34 },
+        { bucket1: 'B', bucket2: 'F', bucket3: 'Z', metric1: 35, metric2: 36 },
+        { bucket1: 'C', bucket2: 'D', bucket3: 'X', metric1: 37, metric2: 38 },
+        { bucket1: 'C', bucket2: 'D', bucket3: 'Y', metric1: 39, metric2: 40 },
+        { bucket1: 'C', bucket2: 'D', bucket3: 'Z', metric1: 41, metric2: 42 },
+        { bucket1: 'C', bucket2: 'E', bucket3: 'X', metric1: 43, metric2: 44 },
+        { bucket1: 'C', bucket2: 'E', bucket3: 'Y', metric1: 45, metric2: 46 },
+        { bucket1: 'C', bucket2: 'E', bucket3: 'Z', metric1: 47, metric2: 48 },
+        { bucket1: 'C', bucket2: 'F', bucket3: 'X', metric1: 49, metric2: 50 },
+        { bucket1: 'C', bucket2: 'F', bucket3: 'Y', metric1: 51, metric2: 52 },
+        { bucket1: 'C', bucket2: 'F', bucket3: 'Z', metric1: 53, metric2: 54 },
+      ],
+    };
+  }
+
+  function buildArgs(): DatatableArgs {
+    return {
+      title: 'Table',
+      sortingColumnId: undefined,
+      sortingDirection: 'none',
+      columns: [
+        {
+          type: 'lens_datatable_column',
+          columnId: 'bucket1',
+          isTransposed: false,
+          transposable: false,
+        },
+        {
+          type: 'lens_datatable_column',
+          columnId: 'bucket2',
+          isTransposed: false,
+          transposable: false,
+        },
+        {
+          type: 'lens_datatable_column',
+          columnId: 'bucket3',
+          isTransposed: false,
+          transposable: false,
+        },
+        {
+          type: 'lens_datatable_column',
+          columnId: 'metric1',
+          isTransposed: false,
+          transposable: true,
+        },
+        {
+          type: 'lens_datatable_column',
+          columnId: 'metric2',
+          isTransposed: false,
+          transposable: true,
+        },
+      ],
+    };
+  }
+
+  it('should correctly sort columns in table by order of args.columns', async () => {
+    const table = buildTable();
+    const shuffledTable: Datatable = {
+      ...table,
+      columns: shuffle(table.columns),
+    };
+    const args = buildArgs();
+    const result = await datatableFn(() => mockFormatFactory)(shuffledTable, args, context);
+
+    const resultColumnIds = result.value.data.columns.map((c) => c.id);
+    const expectedColumnIds = args.columns.map((c) => c.columnId);
+
+    expect(resultColumnIds).toEqual(expectedColumnIds);
+    expect(resultColumnIds).toEqual(['bucket1', 'bucket2', 'bucket3', 'metric1', 'metric2']);
+  });
+});

x-pack/platform/plugins/shared/lens/common/expressions/impl/datatable/datatable_fn.ts

@@ -20,12 +20,19 @@ export const datatableFn =
     getFormatFactory: (context: ExecutionContext) => FormatFactory | Promise<FormatFactory>
   ): DatatableExpressionFunction['fn'] =>
   async (table, args, context) => {
+    const columnSortMap = args.columns.reduce((acc, c, i) => acc.set(c.columnId, i), new Map());
+    const getColumnSort = (id: string) => columnSortMap.get(id) ?? -1;
+    const sortedTable: Datatable = {
+      ...table,
+      columns: table.columns.slice().sort((a, b) => getColumnSort(a.id) - getColumnSort(b.id)),
+    };
+
     if (context?.inspectorAdapters?.tables) {
       context.inspectorAdapters.tables.reset();
       context.inspectorAdapters.tables.allowCsvExport = true;
 
       const logTable = prepareLogTable(
-        table,
+        sortedTable,
         [
           [
             args.columns.map((column) => column.columnId),
@@ -45,20 +52,20 @@ export const datatableFn =
     const formatters: Record<string, ReturnType<FormatFactory>> = {};
     const formatFactory = await getFormatFactory(context);
 
-    table.columns.forEach((column) => {
+    sortedTable.columns.forEach((column) => {
       formatters[column.id] = formatFactory(column.meta?.params);
     });
 
     const hasTransposedColumns = args.columns.some((c) => c.isTransposed);
     if (hasTransposedColumns) {
       // store original shape of data separately
-      untransposedData = cloneDeep(table);
+      untransposedData = cloneDeep(sortedTable);
       // transposes table and args in-place
-      transposeTable(args, table, formatters);
+      transposeTable(args, sortedTable, formatters);
 
       if (context?.inspectorAdapters?.tables) {
         const logTransposedTable = prepareLogTable(
-          table,
+          sortedTable,
           [
             [
               args.columns.map((column) => column.columnId),
@@ -82,7 +89,7 @@ export const datatableFn =
     for (const column of columnsWithSummary) {
       column.summaryRowValue = computeSummaryRowForColumn(
         column,
-        table,
+        sortedTable,
         formatters,
         formatFactory({ id: 'number' })
       );
@@ -92,7 +99,7 @@ export const datatableFn =
       type: 'render',
       as: 'lens_datatable_renderer',
       value: {
-        data: table,
+        data: sortedTable,
         untransposedData,
         syncColors: context.isSyncColorsEnabled?.() ?? false,
         args: {