[Lens] Fix table sorting bug (#74902)

* [Lens] Fix table sorting bug

* Fix types

x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx

@@ -4,6 +4,8 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+import { Ast } from '@kbn/interpreter/common';
+import { buildExpression } from '../../../../../src/plugins/expressions/public';
 import { createMockDatasource } from '../editor_frame_service/mocks';
 import { DatatableVisualizationState, datatableVisualization } from './visualization';
 import { Operation, DataType, FramePublicAPI, TableSuggestionColumn } from '../types';
@@ -324,4 +326,27 @@ describe('Datatable Visualization', () => {
       });
     });
   });
+
+  describe('#toExpression', () => {
+    it('reorders the rendered colums based on the order from the datasource', () => {
+      const datasource = createMockDatasource('test');
+      const layer = { layerId: 'a', columns: ['b', 'c'] };
+      const frame = mockFrame();
+      frame.datasourceLayers = { a: datasource.publicAPIMock };
+      datasource.publicAPIMock.getTableSpec.mockReturnValue([{ columnId: 'c' }, { columnId: 'b' }]);
+      datasource.publicAPIMock.getOperationForColumnId.mockReturnValue({
+        dataType: 'string',
+        isBucketed: true,
+        label: 'label',
+      });
+
+      const expression = datatableVisualization.toExpression({ layers: [layer] }, frame) as Ast;
+      const tableArgs = buildExpression(expression).findFunction('lens_datatable_columns');
+
+      expect(tableArgs).toHaveLength(1);
+      expect(tableArgs[0].arguments).toEqual({
+        columnIds: ['c', 'b'],
+      });
+    });
+  });
 });

x-pack/plugins/lens/public/datatable_visualization/visualization.tsx

@@ -4,6 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+import { Ast } from '@kbn/interpreter/common';
 import { i18n } from '@kbn/i18n';
 import { SuggestionRequest, Visualization, VisualizationSuggestion, Operation } from '../types';
 import chartTableSVG from '../assets/chart_datatable.svg';
@@ -185,10 +186,13 @@ export const datatableVisualization: Visualization<
     };
   },
 
-  toExpression(state, frame) {
+  toExpression(state, frame): Ast {
     const layer = state.layers[0];
     const datasource = frame.datasourceLayers[layer.layerId];
-    const operations = layer.columns
+    const originalOrder = datasource.getTableSpec().map(({ columnId }) => columnId);
+    // When we add a column it could be empty, and therefore have no order
+    const sortedColumns = Array.from(new Set(originalOrder.concat(layer.columns)));
+    const operations = sortedColumns
       .map((columnId) => ({ columnId, operation: datasource.getOperationForColumnId(columnId) }))
       .filter((o): o is { columnId: string; operation: Operation } => !!o.operation);