fix(partition): consider legend extras when computing the legend size

packages/charts/src/chart_types/partition_chart/layout/utils/legend_labels.ts

@@ -7,6 +7,7 @@
  */
 
 import { LegendItemLabel } from '../../../../state/selectors/get_legend_items_labels';
+import { ValueFormatter } from '../../../../utils/common';
 import { Layer } from '../../specs';
 import {
   CHILDREN_KEY,
@@ -15,15 +16,22 @@ import {
   PATH_KEY,
   ArrayNode,
   NULL_SMALL_MULTIPLES_KEY,
+  AGGREGATE_KEY,
 } from './group_by_rollup';
 
 /** @internal */
-export function getLegendLabels(layers: Layer[], tree: HierarchyOfArrays, legendMaxDepth: number) {
-  return flatSlicesNames(layers, 0, tree).filter(({ depth }) => depth <= legendMaxDepth);
+export function getLegendLabelsAndValue(
+  layers: Layer[],
+  tree: HierarchyOfArrays,
+  legendMaxDepth: number,
+  // if used the resulting label will be concatenated with the formatted value, use () => '' to avoid that
+  valueFormatter: ValueFormatter,
+): LegendItemLabel[] {
+  return flatSlicesNames(layers, 0, tree, valueFormatter).filter(({ depth }) => depth <= legendMaxDepth);
 }
 
 /** @internal */
-export function getArrayNodeKey(arrayNode: ArrayNode): string {
+function getArrayNodeKey(arrayNode: ArrayNode): string {
   return arrayNode[PATH_KEY].reduce<string>((acc, { value, index }) => {
     if (value === HIERARCHY_ROOT_KEY || value === NULL_SMALL_MULTIPLES_KEY) return acc;
     return `${acc}(${index}):${value}__`;
@@ -34,6 +42,7 @@ function flatSlicesNames(
   layers: Layer[],
   depth: number,
   tree: HierarchyOfArrays,
+  valueFormatter: ValueFormatter,
   keys: Map<string, string> = new Map(),
   depths: Map<string, number> = new Map(),
 ): LegendItemLabel[] {
@@ -45,18 +54,20 @@ function flatSlicesNames(
     // format the key with the layer formatter
     const layer = layers[depth - 1];
     const formatter = layer?.nodeLabel;
-    const formattedValue = formatter ? formatter(key) : `${key}`;
+    const formattedKey = formatter ? formatter(key) : `${key}`;
     // preventing errors from external formatters
-    if (formattedValue && formattedValue !== HIERARCHY_ROOT_KEY) {
+    if (formattedKey && formattedKey !== HIERARCHY_ROOT_KEY) {
       // Node key must be unique for each node in the tree
       const nodeKey = getArrayNodeKey(arrayNode);
       // save only the max depth, so we can compute the the max extension of the legend
       depths.set(nodeKey, depth);
-      keys.set(nodeKey, formattedValue);
+
+      const formattedValue = valueFormatter(arrayNode[AGGREGATE_KEY]);
+      keys.set(nodeKey, `${formattedKey}${formattedValue}`);
     }
 
     const children = arrayNode[CHILDREN_KEY];
-    flatSlicesNames(layers, depth + 1, children, keys, depths);
+    flatSlicesNames(layers, depth + 1, children, valueFormatter, keys, depths);
   }
 
   return [...depths.keys()].map((key) => ({

packages/charts/src/chart_types/partition_chart/state/chart_state.tsx

@@ -61,6 +61,8 @@ export class PartitionState implements InternalChartState {
 
   getLegendItemsLabels(globalState: GlobalChartState) {
     // order doesn't matter, but it needs to return the highest depth of the label occurrence so enough horizontal width is allocated
+    // the label item strings needs to be a concatenation of the label + the extra formatted value if available.
+    // this is required to compute the legend automatic width
     return getLegendItemsLabels(globalState);
   }
 

packages/charts/src/chart_types/partition_chart/state/selectors/get_legend_items_labels.ts

@@ -9,15 +9,19 @@
 import { createCustomCachedSelector } from '../../../../state/create_selector';
 import { LegendItemLabel } from '../../../../state/selectors/get_legend_items_labels';
 import { getSettingsSpecSelector } from '../../../../state/selectors/get_settings_specs';
-import { getLegendLabels } from '../../layout/utils/legend_labels';
+import { getLegendLabelsAndValue } from '../../layout/utils/legend_labels';
 import { getPartitionSpecs } from './get_partition_specs';
 import { getTrees } from './tree';
 
 /** @internal */
 export const getLegendItemsLabels = createCustomCachedSelector(
   [getPartitionSpecs, getSettingsSpecSelector, getTrees],
-  (specs, { legendMaxDepth, showLegend }, trees): LegendItemLabel[] =>
-    specs.flatMap(({ layers }) =>
-      showLegend ? trees.flatMap(({ tree }) => getLegendLabels(layers, tree, legendMaxDepth)) : [],
+  (specs, { legendMaxDepth, showLegend, showLegendExtra }, trees): LegendItemLabel[] =>
+    specs.flatMap(({ layers, valueFormatter }) =>
+      showLegend
+        ? trees.flatMap(({ tree }) =>
+            getLegendLabelsAndValue(layers, tree, legendMaxDepth, showLegendExtra ? valueFormatter : () => ''),
+          )
+        : [],
     ),
 );

storybook/stories/waffle/1_simple.story.tsx

@@ -22,13 +22,14 @@ import { mocks } from '@elastic/charts/src/mocks/hierarchical';
 
 import { ArrayEntry } from '../../../packages/charts/src/chart_types/partition_chart/layout/utils/group_by_rollup';
 import { useBaseTheme } from '../../use_base_theme';
-import { colorBrewerCategoricalStark9, discreteColor, productLookup } from '../utils/utils';
+import { colorBrewerCategoricalStark9, discreteColor } from '../utils/utils';
 
 export const Example = () => {
   const showDebug = boolean('show table for debugging', false);
   const ascendingSort = boolean('ascending sort', false);
   // this is used to test the sorting capabilities
   const data = mocks.pie.slice(0, 4).sort(() => (Math.random() > 0.5 ? 1 : -1));
+  const names: Record<string, string> = { '7': 'Al', '3': 'Au', '5': 'Ag', '8': 'Cu' };
   return (
     <Chart className="story-chart">
       <Settings
@@ -51,7 +52,7 @@ export const Example = () => {
         layers={[
           {
             groupByRollup: (d: Datum) => d.sitc1,
-            nodeLabel: (d: Datum) => productLookup[d].name,
+            nodeLabel: (d: Datum) => names[d],
             shape: {
               fillColor: (d: ShapeTreeNode) => discreteColor(colorBrewerCategoricalStark9.slice(1))(d.sortIndex),
             },