[Lens] Improve pie suggestions (#102755) (#103461)

Co-authored-by: Joe Reuter <johannes.reuter@elastic.co>
elastic · Jun 28, 2021 · 4863816 · 4863816
1 parent 67ce486
commit 4863816
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 6 deletions.
diff --git a/x-pack/plugins/lens/public/pie_visualization/suggestions.test.ts b/x-pack/plugins/lens/public/pie_visualization/suggestions.test.ts
@@ -6,7 +6,7 @@
  */
 
 import { PaletteOutput } from 'src/plugins/charts/public';
-import { DataType } from '../types';
+import { DataType, SuggestionRequest } from '../types';
 import { suggestions } from './suggestions';
 import { PieVisualizationState } from './types';
 
@@ -354,6 +354,81 @@ describe('suggestions', () => {
       );
     });
 
+    it('should score higher for more groups', () => {
+      const config: SuggestionRequest<PieVisualizationState> = {
+        table: {
+          layerId: 'first',
+          isMultiRow: true,
+          columns: [
+            {
+              columnId: 'a',
+              operation: { label: 'Top 5', dataType: 'string' as DataType, isBucketed: true },
+            },
+            {
+              columnId: 'b',
+              operation: { label: 'Top 5', dataType: 'string' as DataType, isBucketed: true },
+            },
+            {
+              columnId: 'e',
+              operation: { label: 'Count', dataType: 'number' as DataType, isBucketed: false },
+            },
+          ],
+          changeType: 'initial',
+        },
+        state: undefined,
+        keptLayerIds: ['first'],
+      };
+      const twoGroupsResults = suggestions(config);
+      config.table.columns.splice(1, 1);
+      const oneGroupResults = suggestions(config);
+
+      expect(Math.max(...twoGroupsResults.map((suggestion) => suggestion.score))).toBeGreaterThan(
+        Math.max(...oneGroupResults.map((suggestion) => suggestion.score))
+      );
+    });
+
+    it('should score higher for more groups for each subvis with passed-in subvis id', () => {
+      const config: SuggestionRequest<PieVisualizationState> = {
+        table: {
+          layerId: 'first',
+          isMultiRow: true,
+          columns: [
+            {
+              columnId: 'a',
+              operation: { label: 'Top 5', dataType: 'string' as DataType, isBucketed: true },
+            },
+            {
+              columnId: 'b',
+              operation: { label: 'Top 5', dataType: 'string' as DataType, isBucketed: true },
+            },
+            {
+              columnId: 'e',
+              operation: { label: 'Count', dataType: 'number' as DataType, isBucketed: false },
+            },
+          ],
+          changeType: 'initial',
+        },
+        state: undefined,
+        keptLayerIds: ['first'],
+        subVisualizationId: 'donut',
+      };
+      const twoGroupsResults = suggestions(config);
+      config.table.columns.splice(1, 1);
+      const oneGroupResults = suggestions(config);
+      // collect scores for one or two groups for each sub vis
+      const scores: Record<string, { two: number; one: number }> = {};
+      twoGroupsResults.forEach((r) => {
+        scores[r.state.shape] = { ...(scores[r.state.shape] || {}), two: r.score };
+      });
+      oneGroupResults.forEach((r) => {
+        scores[r.state.shape] = { ...(scores[r.state.shape] || {}), one: r.score };
+      });
+      expect(Object.keys(scores).length).toEqual(2);
+      Object.values(scores).forEach(({ one, two }) => {
+        expect(two).toBeGreaterThan(one);
+      });
+    });
+
     it('should keep passed in palette', () => {
       const mainPalette: PaletteOutput = { type: 'palette', name: 'mock' };
       const results = suggestions({

diff --git a/x-pack/plugins/lens/public/pie_visualization/suggestions.ts b/x-pack/plugins/lens/public/pie_visualization/suggestions.ts
@@ -51,9 +51,10 @@ export function suggestions({
 
   const results: Array<VisualizationSuggestion<PieVisualizationState>> = [];
 
-  if (groups.length <= MAX_PIE_BUCKETS) {
-    let newShape: PieVisualizationState['shape'] = 'donut';
-    if (groups.length !== 1) {
+  if (groups.length <= MAX_PIE_BUCKETS && subVisualizationId !== 'treemap') {
+    let newShape: PieVisualizationState['shape'] =
+      (subVisualizationId as PieVisualizationState['shape']) || 'donut';
+    if (groups.length !== 1 && !subVisualizationId) {
       newShape = 'pie';
     }
 
@@ -108,7 +109,10 @@ export function suggestions({
     });
   }
 
-  if (groups.length <= MAX_TREEMAP_BUCKETS) {
+  if (
+    groups.length <= MAX_TREEMAP_BUCKETS &&
+    (!subVisualizationId || subVisualizationId === 'treemap')
+  ) {
     results.push({
       title: i18n.translate('xpack.lens.pie.treemapSuggestionLabel', {
         defaultMessage: 'As Treemap',
@@ -149,7 +153,11 @@ export function suggestions({
   }
 
   return [...results]
-    .sort((a, b) => a.score - b.score)
+    .map((suggestion) => ({
+      ...suggestion,
+      score: suggestion.score + 0.05 * groups.length,
+    }))
+    .sort((a, b) => b.score - a.score)
     .map((suggestion) => ({
       ...suggestion,
       hide: incompleteConfiguration || suggestion.hide,