fix(partition): allow custom sorting for the legend items (#1959)

The legendSort prop is now applied to partition charts allowing the user to custom sort (flattened) legend items.

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

@@ -9,7 +9,7 @@
 import { CategoryKey } from '../../../../common/category';
 import { map } from '../../../../common/iterables';
 import { LegendItem } from '../../../../common/legend';
-import { LegendPositionConfig } from '../../../../specs/settings';
+import { LegendPositionConfig, SettingsSpec } from '../../../../specs/settings';
 import { isHierarchicalLegend } from '../../../../utils/legend';
 import { Layer } from '../../specs';
 import { PartitionLayout } from '../types/config_types';
@@ -34,6 +34,10 @@ function compareTreePaths(
   return a.length - b.length; // if one path is fully contained in the other, then parent (shorter) goes first
 }
 
+function createPartitionIdentifier(id: string, item: QuadViewModel) {
+  return { specId: id, key: item.dataName, smAccessorValue: item.smAccessorValue };
+}
+
 /** @internal */
 export function getLegendItems(
   id: string,
@@ -43,9 +47,16 @@ export function getLegendItems(
   legendPosition: LegendPositionConfig,
   quadViewModel: QuadViewModel[],
   partitionLayout: PartitionLayout | undefined,
+  settingsSpec: SettingsSpec,
 ): LegendItem[] {
   const uniqueNames = new Set(map(quadViewModel, ({ dataName, fillColor }) => makeKey(dataName, fillColor)));
   const useHierarchicalLegend = isHierarchicalLegend(flatLegend, legendPosition);
+  const sortingFn = flatLegend && settingsSpec.legendSort;
+
+  const customSortingFn = sortingFn
+    ? (aItem: QuadViewModel, bItem: QuadViewModel) =>
+        sortingFn(createPartitionIdentifier(id, aItem), createPartitionIdentifier(id, bItem))
+    : null;
 
   const formattedLabel = ({ dataName, depth }: QuadViewModel) => {
     const formatter = layers[depth - 1]?.nodeLabel;
@@ -77,11 +88,14 @@ export function getLegendItems(
     return true;
   });
 
+  const waffleSortingFn = customSortingFn ?? descendingValues;
+  const flatLegendSortingFn = customSortingFn ?? compareNames;
+
   items.sort(
     partitionLayout === PartitionLayout.waffle // waffle has inherent top to bottom descending order
-      ? descendingValues
+      ? waffleSortingFn
       : flatLegend
-      ? compareNames
+      ? flatLegendSortingFn
       : compareTreePaths,
   );
 

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

@@ -11,12 +11,13 @@ import { getPartitionSpecs } from './get_partition_specs';
 import { LegendItem } from '../../../../common/legend';
 import { createCustomCachedSelector } from '../../../../state/create_selector';
 import { getLegendConfigSelector } from '../../../../state/selectors/get_legend_config_selector';
+import { getSettingsSpecSelector } from '../../../../state/selectors/get_settings_spec';
 import { getLegendItems } from '../../layout/utils/legend';
 
 /** @internal */
 export const computeLegendSelector = createCustomCachedSelector(
-  [getPartitionSpecs, getLegendConfigSelector, partitionMultiGeometries],
-  (specs, { flatLegend, legendMaxDepth, legendPosition }, geometries): LegendItem[] =>
+  [getPartitionSpecs, getLegendConfigSelector, partitionMultiGeometries, getSettingsSpecSelector],
+  (specs, { flatLegend, legendMaxDepth, legendPosition }, geometries, settings): LegendItem[] =>
     specs.flatMap((partitionSpec, i) => {
       const quadViewModel = geometries.filter((g) => g.index === i).flatMap((g) => g.quadViewModel);
       return getLegendItems(
@@ -27,6 +28,7 @@ export const computeLegendSelector = createCustomCachedSelector(
         legendPosition,
         quadViewModel,
         partitionSpec.layout,
+        settings,
       );
     }),
 );

storybook/stories/legend/10_sunburst.story.tsx

@@ -19,6 +19,7 @@ import {
   PartitionLayout,
   Settings,
   defaultPartitionValueFormatter,
+  SeriesIdentifier,
 } from '@elastic/charts';
 import { ShapeTreeNode } from '@elastic/charts/src/chart_types/partition_chart/layout/types/viewmodel_types';
 import { mocks } from '@elastic/charts/src/mocks/hierarchical';
@@ -38,6 +39,8 @@ export const Example = () => {
     {
       treemap: PartitionLayout.treemap,
       sunburst: PartitionLayout.sunburst,
+      mosaic: PartitionLayout.mosaic,
+      waffle: PartitionLayout.waffle,
     },
     PartitionLayout.sunburst,
   );
@@ -50,6 +53,26 @@ export const Example = () => {
   });
   const legendStrategy = select('legendStrategy', LegendStrategy, LegendStrategy.Key as LegendStrategy);
   const maxLines = number('max legend label lines', 1, { min: 0, step: 1 });
+
+  const legendSortStrategy = select(
+    'Custom legend sorting',
+    { RegionsFirst: 'regionsFirst', ProductsFirst: 'productsFirst', DefaultSort: 'default' },
+    'regionsFirst',
+  );
+
+  const customLegendSort = (a: SeriesIdentifier, b: SeriesIdentifier) => {
+    if (legendSortStrategy === 'regionsFirst') {
+      if (a.key in regionLookup && b.key in regionLookup) {
+        return a.key.localeCompare(b.key);
+      }
+      return a.key in regionLookup ? -1 : b.key in regionLookup ? 1 : a.key.localeCompare(b.key);
+    }
+    if (a.key in productLookup && b.key in productLookup) {
+      return a.key.localeCompare(b.key);
+    }
+    return a.key in productLookup ? -1 : b.key in productLookup ? 1 : a.key.localeCompare(b.key);
+  };
+
   const partitionTheme: PartialTheme['partition'] = {
     linkLabel: {
       maxCount: 0,
@@ -72,14 +95,18 @@ export const Example = () => {
     circlePadding: 4,
   };
 
+  const isFlatLegendSupported =
+    partitionLayout === PartitionLayout.treemap || partitionLayout === PartitionLayout.sunburst;
+
   return (
     <Chart>
       <Settings
         showLegend
         showLegendExtra={showLegendExtra}
-        flatLegend={flatLegend}
+        flatLegend={isFlatLegendSupported ? flatLegend : true}
         legendStrategy={legendStrategy}
         legendMaxDepth={legendMaxDepth}
+        legendSort={legendSortStrategy !== 'default' ? customLegendSort : undefined}
         baseTheme={useBaseTheme()}
         theme={{
           partition: partitionTheme,
@@ -127,5 +154,8 @@ Example.parameters = {
   markdown: `To flatten a hierarchical legend (like the rendered in a pie chart or a treemap when using a multi-layer configuration) you can
 add the \`flatLegend\` prop into the \`<Settings  baseTheme={useBaseTheme()} />\` component.
 
-To limit displayed hierarchy to a specific depth, you can use the \`legendMaxDepth\` prop. The first layer will have a depth of \`1\`.`,
+To limit displayed hierarchy to a specific depth, you can use the \`legendMaxDepth\` prop. The first layer will have a depth of \`1\`.
+
+It is possible to provide a custom sorting logic for a legend when flattened, when not flattened the \`legendSort\` function is ignored.
+`,
 };