From ac0495d5b136794e8c1e0b98d19571562d48b196 Mon Sep 17 00:00:00 2001
From: Wylie Conlon <wylieconlon@gmail.com>
Date: Wed, 12 Aug 2020 16:27:41 -0400
Subject: [PATCH] [Lens] Fix table sorting bug

---
 .../visualization.test.tsx                    | 24 +++++++++++++++++++
 .../datatable_visualization/visualization.tsx |  5 +++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx b/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx
index e18190b6c2d69..933649126beaa 100644
--- a/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx
+++ b/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx
@@ -4,6 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+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 +325,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);
+      const tableArgs = buildExpression(expression).findFunction('lens_datatable_columns');
+
+      expect(tableArgs).toHaveLength(1);
+      expect(tableArgs[0].arguments).toEqual({
+        columnIds: ['c', 'b'],
+      });
+    });
+  });
 });
diff --git a/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx b/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx
index e4b371143594a..123fb972ad8b1 100644
--- a/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx
+++ b/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx
@@ -188,7 +188,10 @@ export const datatableVisualization: Visualization<
   toExpression(state, frame) {
     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);