[Lens] Partition legend order synced with filters order

src/plugins/chart_expressions/expression_partition_vis/public/components/partition_vis_component.tsx

@@ -413,7 +413,7 @@ const PartitionVisComponent = (props: PartitionVisComponentProps) => {
     [visData.rows, metricColumn]
   );
 
-  const flatLegend = isLegendFlat(visType, splitChartDimension);
+  const flatLegend = !visParams.nestedLegend || isLegendFlat(visType, splitChartDimension);
 
   const canShowPieChart = !isEmpty && !isMetricEmpty && !isAllZeros && !hasNegative;
 
@@ -425,6 +425,32 @@ const PartitionVisComponent = (props: PartitionVisComponentProps) => {
 
   const partitionType = getPartitionType(visType);
 
+  const customLegendSort = useMemo(() => {
+    if (!showLegend || !flatLegend) {
+      return;
+    }
+    const [bucketColumn] = bucketColumns;
+    if (!bucketColumn.id) {
+      return;
+    }
+    const lookup: Record<string, number> = {};
+    visData.rows.forEach((row, i) => {
+      const category = row[bucketColumn.id!];
+      if (!(category in lookup)) {
+        lookup[category] = i;
+      }
+    });
+    return (a: SeriesIdentifier, b: SeriesIdentifier) => {
+      if (a.key == null) {
+        return 1;
+      }
+      if (b.key == null) {
+        return -1;
+      }
+      return lookup[a.key] - lookup[b.key];
+    };
+  }, [bucketColumns, flatLegend, showLegend, visData.rows]);
+
   return (
     <div css={chartContainerStyle} data-test-subj="partitionVisChart">
       {!canShowPieChart ? (
@@ -471,6 +497,7 @@ const PartitionVisComponent = (props: PartitionVisComponentProps) => {
                 legendMaxDepth={visParams.nestedLegend ? undefined : 1}
                 legendColorPicker={props.uiState ? LegendColorPickerWrapper : undefined}
                 flatLegend={flatLegend}
+                legendSort={customLegendSort}
                 tooltip={tooltip}
                 showLegendExtra={visParams.showValuesInLegend}
                 onElementClick={([elementEvent]) => {

x-pack/plugins/lens/public/shared_components/legend/legend_settings_popover.tsx

@@ -165,30 +165,32 @@ export const MaxLinesInput = ({
   );
 };
 
+const noop = () => {};
+
 export const LegendSettingsPopover: React.FunctionComponent<LegendSettingsPopoverProps> = ({
   legendOptions,
   mode,
   onDisplayChange,
   position,
   location,
-  onLocationChange = () => {},
+  onLocationChange = noop,
   verticalAlignment,
   horizontalAlignment,
   floatingColumns,
-  onAlignmentChange = () => {},
-  onFloatingColumnsChange = () => {},
+  onAlignmentChange = noop,
+  onFloatingColumnsChange = noop,
   onPositionChange,
   renderNestedLegendSwitch,
   nestedLegend,
-  onNestedLegendChange = () => {},
+  onNestedLegendChange = noop,
   valueInLegend,
-  onValueInLegendChange = () => {},
+  onValueInLegendChange = noop,
   renderValueInLegendSwitch,
   groupPosition = 'right',
   maxLines,
-  onMaxLinesChange = () => {},
+  onMaxLinesChange = noop,
   shouldTruncate,
-  onTruncateLegendChange = () => {},
+  onTruncateLegendChange = noop,
   legendSize,
   onLegendSizeChange,
   showAutoLegendSizeOption,
@@ -294,7 +296,7 @@ export const LegendSettingsPopover: React.FunctionComponent<LegendSettingsPopove
                 })}
                 data-test-subj="lens-legend-nested-switch"
                 showLabel={false}
-                checked={!!nestedLegend}
+                checked={Boolean(nestedLegend)}
                 onChange={onNestedLegendChange}
               />
             </EuiFormRow>

x-pack/plugins/lens/public/visualizations/partition/to_expression.ts

@@ -192,7 +192,10 @@ const generateCommonArguments = (
     legendPosition: layer.legendPosition || Position.Right,
     maxLegendLines: layer.legendMaxLines ?? 1,
     legendSize: layer.legendSize,
-    nestedLegend: !!layer.nestedLegend,
+    nestedLegend:
+      layer.primaryGroups.length + (layer.secondaryGroups?.length ?? 0) > 1
+        ? Boolean(layer.nestedLegend)
+        : false,
     truncateLegend:
       layer.truncateLegend ?? getDefaultVisualValuesForLayer(state, datasourceLayers).truncateText,
     palette: generatePaletteAstArguments(paletteService, state.palette),

x-pack/plugins/lens/public/visualizations/partition/toolbar.tsx

@@ -263,7 +263,10 @@ export function PieToolbar(props: VisualizationToolbarProps<PieVisualizationStat
         onValueInLegendChange={onValueInLegendChange}
         position={layer.legendPosition}
         onPositionChange={onLegendPositionChange}
-        renderNestedLegendSwitch={!PartitionChartsMeta[state.shape]?.legend.hideNestedLegendSwitch}
+        renderNestedLegendSwitch={
+          !PartitionChartsMeta[state.shape]?.legend.hideNestedLegendSwitch &&
+          layer.primaryGroups.length + (layer.secondaryGroups?.length ?? 0) > 1
+        }
         nestedLegend={Boolean(layer.nestedLegend)}
         onNestedLegendChange={onNestedLegendChange}
         shouldTruncate={layer.truncateLegend ?? defaultTruncationValue}