[Lens] Suggest reduced versions of current data table

x-pack/legacy/plugins/lens/public/datatable_visualization_plugin/visualization.tsx

@@ -130,6 +130,7 @@ export const datatableVisualization: Visualization<
   }: SuggestionRequest<DatatableVisualizationState>): Array<
     VisualizationSuggestion<DatatableVisualizationState>
   > {
+    const maxColumnCount = Math.max.apply(undefined, tables.map(table => table.columns.length));
     return tables.map(table => {
       const title = i18n.translate('xpack.lens.datatable.visualizationOf', {
         defaultMessage: 'Table: {operations}',
@@ -148,7 +149,8 @@ export const datatableVisualization: Visualization<
 
       return {
         title,
-        score: 1,
+        // largest possible table will have a score of 0.2, less columns reduce score
+        score: (table.columns.length / maxColumnCount) * 0.2,
         datasourceSuggestionId: table.datasourceSuggestionId,
         state: {
           layers: [

x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/__mocks__/suggestion_helpers.ts

@@ -8,4 +8,4 @@ const actual = jest.requireActual('../suggestion_helpers');
 
 jest.spyOn(actual, 'getSuggestions');
 
-export const { getSuggestions, toSwitchAction } = actual;
+export const { getSuggestions, switchToSuggestion } = actual;

x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/chart_switch.test.tsx

@@ -5,12 +5,13 @@
  */
 
 import React from 'react';
-import { createMockVisualization, createMockFramePublicAPI } from '../mocks';
+import { createMockVisualization, createMockFramePublicAPI, createMockDatasource } from '../mocks';
 import { mountWithIntl as mount } from 'test_utils/enzyme_helpers';
 import { ReactWrapper } from 'enzyme';
 import { ChartSwitch } from './chart_switch';
 import { Visualization, FramePublicAPI, DatasourcePublicAPI } from '../../types';
 import { EuiKeyPadMenuItemButton } from '@elastic/eui';
+import { Action } from './state_management';
 
 describe('chart_switch', () => {
   function generateVisualization(id: string): jest.Mocked<Visualization> {
@@ -83,6 +84,33 @@ describe('chart_switch', () => {
     } as FramePublicAPI;
   }
 
+  function mockDatasourceMap() {
+    const datasource = createMockDatasource();
+    datasource.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      {
+        state: {},
+        table: {
+          columns: [],
+          datasourceSuggestionId: 0,
+          isMultiRow: true,
+          layerId: 'a',
+        },
+      },
+    ]);
+    return {
+      testDatasource: datasource,
+    };
+  }
+
+  function mockDatasourceStates() {
+    return {
+      testDatasource: {
+        state: {},
+        isLoading: false,
+      },
+    };
+  }
+
   function showFlyout(component: ReactWrapper) {
     component
       .find('[data-test-subj="lnsChartSwitchPopover"]')
@@ -116,6 +144,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={mockFrame(['a'])}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -125,6 +155,8 @@ describe('chart_switch', () => {
       initialState: 'suggestion visB',
       newVisualizationId: 'visB',
       type: 'SWITCH_VISUALIZATION',
+      datasourceId: 'testDatasource',
+      datasourceState: {},
     });
   });
 
@@ -142,6 +174,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -154,6 +188,61 @@ describe('chart_switch', () => {
     });
   });
 
+  it('should indicate data loss if not all columns will be used', () => {
+    const dispatch = jest.fn();
+    const visualizations = mockVisualizations();
+    const frame = mockFrame(['a']);
+
+    const datasourceMap = mockDatasourceMap();
+    datasourceMap.testDatasource.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      {
+        state: {},
+        table: {
+          columns: [
+            {
+              columnId: 'col1',
+              operation: {
+                label: '',
+                dataType: 'string',
+                isBucketed: true,
+              },
+            },
+            {
+              columnId: 'col2',
+              operation: {
+                label: '',
+                dataType: 'number',
+                isBucketed: false,
+              },
+            },
+          ],
+          datasourceSuggestionId: 0,
+          layerId: 'first',
+          isMultiRow: true,
+        },
+      },
+    ]);
+    datasourceMap.testDatasource.publicAPIMock.getTableSpec.mockReturnValue([
+      { columnId: 'col1' },
+      { columnId: 'col2' },
+      { columnId: 'col3' },
+    ]);
+
+    const component = mount(
+      <ChartSwitch
+        visualizationId="visA"
+        visualizationState={{}}
+        visualizationMap={visualizations}
+        dispatch={dispatch}
+        framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
+      />
+    );
+
+    expect(getMenuItem('subvisB', component).prop('betaBadgeIconType')).toEqual('alert');
+  });
+
   it('should indicate data loss if not all layers will be used', () => {
     const dispatch = jest.fn();
     const visualizations = mockVisualizations();
@@ -166,6 +255,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -185,6 +276,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -205,6 +298,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -230,6 +325,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -249,6 +346,8 @@ describe('chart_switch', () => {
         visualizationMap={mockVisualizations()}
         dispatch={jest.fn()}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -258,6 +357,30 @@ describe('chart_switch', () => {
     expect(removeLayers).toHaveBeenCalledWith(['b', 'c']);
   });
 
+  it('should remove all layers if there is no suggestion', () => {
+    const dispatch = jest.fn();
+    const visualizations = mockVisualizations();
+    visualizations.visB.getSuggestions.mockReturnValueOnce([]);
+    const frame = mockFrame(['a', 'b', 'c']);
+
+    const component = mount(
+      <ChartSwitch
+        visualizationId="visA"
+        visualizationState={{}}
+        visualizationMap={visualizations}
+        dispatch={dispatch}
+        framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
+      />
+    );
+
+    switchTo('subvisB', component);
+
+    expect(frame.removeLayers).toHaveBeenCalledTimes(1);
+    expect(frame.removeLayers).toHaveBeenCalledWith(['a', 'b', 'c']);
+  });
+
   it('should not remove layers if the visualization is not changing', () => {
     const dispatch = jest.fn();
     const removeLayers = jest.fn();
@@ -277,16 +400,78 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
     switchTo('subvisC2', component);
     expect(removeLayers).not.toHaveBeenCalled();
     expect(switchVisualizationType).toHaveBeenCalledWith('subvisC2', 'therebegriffins');
+    expect(dispatch).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: 'SWITCH_VISUALIZATION',
+        initialState: 'therebedragons',
+      })
+    );
+  });
+
+  it('should switch to the updated datasource state', () => {
+    const dispatch = jest.fn();
+    const visualizations = mockVisualizations();
+    const frame = mockFrame(['a', 'b']);
+
+    const datasourceMap = mockDatasourceMap();
+    datasourceMap.testDatasource.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      {
+        state: 'testDatasource suggestion',
+        table: {
+          columns: [
+            {
+              columnId: 'col1',
+              operation: {
+                label: '',
+                dataType: 'string',
+                isBucketed: true,
+              },
+            },
+            {
+              columnId: 'col2',
+              operation: {
+                label: '',
+                dataType: 'number',
+                isBucketed: false,
+              },
+            },
+          ],
+          datasourceSuggestionId: 0,
+          layerId: 'a',
+          isMultiRow: true,
+        },
+      },
+    ]);
+
+    const component = mount(
+      <ChartSwitch
+        visualizationId="visA"
+        visualizationState={{}}
+        visualizationMap={visualizations}
+        dispatch={dispatch}
+        framePublicAPI={frame}
+        datasourceMap={datasourceMap}
+        datasourceStates={mockDatasourceStates()}
+      />
+    );
+
+    switchTo('subvisB', component);
+
     expect(dispatch).toHaveBeenCalledWith({
-      type: 'UPDATE_VISUALIZATION_STATE',
-      newState: 'therebedragons',
-    });
+      type: 'SWITCH_VISUALIZATION',
+      newVisualizationId: 'visB',
+      datasourceId: 'testDatasource',
+      datasourceState: 'testDatasource suggestion',
+      initialState: 'suggestion visB',
+    } as Action);
   });
 
   it('should ensure the new visualization has the proper subtype', () => {
@@ -305,6 +490,8 @@ describe('chart_switch', () => {
         visualizationMap={visualizations}
         dispatch={dispatch}
         framePublicAPI={mockFrame(['a'])}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );
 
@@ -314,6 +501,8 @@ describe('chart_switch', () => {
       initialState: 'suggestion visB subvisB',
       newVisualizationId: 'visB',
       type: 'SWITCH_VISUALIZATION',
+      datasourceId: 'testDatasource',
+      datasourceState: {},
     });
   });
 
@@ -325,6 +514,8 @@ describe('chart_switch', () => {
         visualizationMap={mockVisualizations()}
         dispatch={jest.fn()}
         framePublicAPI={mockFrame(['a', 'b'])}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
       />
     );