[Lens] Fixes mosaic with 2 axis coloring (#167035)

## Summary

Fixes https://github.com/elastic/kibana/issues/164964

The Other label was not formatted correctly.

<img width="1226" alt="image"
src="a8c8272e-b0c9-4088-93fc-45432d084570">


### Checklist
- [ ] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

src/plugins/chart_expressions/expression_partition_vis/public/utils/layers/get_color.test.ts

@@ -8,7 +8,6 @@
 
 import type { PaletteOutput, PaletteDefinition } from '@kbn/coloring';
 import { chartPluginMock } from '@kbn/charts-plugin/public/mocks';
-import { Datatable } from '@kbn/expressions-plugin/common';
 import { byDataColorPaletteMap, SimplifiedArrayNode } from './get_color';
 import type { SeriesLayer } from '@kbn/coloring';
 import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
@@ -21,29 +20,28 @@ import { ChartTypes } from '../../../common/types';
 import { getDistinctSeries } from '..';
 
 describe('#byDataColorPaletteMap', () => {
-  let datatable: Datatable;
   let paletteDefinition: PaletteDefinition;
   let palette: PaletteOutput;
-  const columnId = 'foo';
+  const visData = createMockVisData();
+  const defaultFormatter = jest.fn((...args) => fieldFormatsMock.deserialize(...args));
+  const formatters = generateFormatters(visData, defaultFormatter);
 
   beforeEach(() => {
-    datatable = {
-      rows: [
-        {
-          [columnId]: '1',
-        },
-        {
-          [columnId]: '2',
-        },
-      ],
-    } as unknown as Datatable;
     paletteDefinition = chartPluginMock.createPaletteRegistry().get('default');
     palette = { type: 'palette' } as PaletteOutput;
   });
 
   it('should create byDataColorPaletteMap', () => {
-    expect(byDataColorPaletteMap(datatable.rows, columnId, paletteDefinition, palette))
-      .toMatchInlineSnapshot(`
+    expect(
+      byDataColorPaletteMap(
+        visData.rows,
+        visData.columns[0],
+        paletteDefinition,
+        palette,
+        formatters,
+        fieldFormatsMock
+      )
+    ).toMatchInlineSnapshot(`
       Object {
         "getColor": [Function],
       }
@@ -52,21 +50,25 @@ describe('#byDataColorPaletteMap', () => {
 
   it('should get color', () => {
     const colorPaletteMap = byDataColorPaletteMap(
-      datatable.rows,
-      columnId,
+      visData.rows,
+      visData.columns[0],
       paletteDefinition,
-      palette
+      palette,
+      formatters,
+      fieldFormatsMock
     );
 
-    expect(colorPaletteMap.getColor('1')).toBe('black');
+    expect(colorPaletteMap.getColor('Logstash Airways')).toBe('black');
   });
 
   it('should return undefined in case if values not in datatable', () => {
     const colorPaletteMap = byDataColorPaletteMap(
-      datatable.rows,
-      columnId,
+      visData.rows,
+      visData.columns[0],
       paletteDefinition,
-      palette
+      palette,
+      formatters,
+      fieldFormatsMock
     );
 
     expect(colorPaletteMap.getColor('wrong')).toBeUndefined();
@@ -74,24 +76,26 @@ describe('#byDataColorPaletteMap', () => {
 
   it('should increase rankAtDepth for each new value', () => {
     const colorPaletteMap = byDataColorPaletteMap(
-      datatable.rows,
-      columnId,
+      visData.rows,
+      visData.columns[0],
       paletteDefinition,
-      palette
+      palette,
+      formatters,
+      fieldFormatsMock
     );
-    colorPaletteMap.getColor('1');
-    colorPaletteMap.getColor('2');
+    colorPaletteMap.getColor('Logstash Airways');
+    colorPaletteMap.getColor('JetBeats');
 
     expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith(
       1,
-      [{ name: '1', rankAtDepth: 0, totalSeriesAtDepth: 2 }],
+      [{ name: 'Logstash Airways', rankAtDepth: 0, totalSeriesAtDepth: 4 }],
       { behindText: false },
       undefined
     );
 
     expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith(
       2,
-      [{ name: '2', rankAtDepth: 1, totalSeriesAtDepth: 2 }],
+      [{ name: 'JetBeats', rankAtDepth: 1, totalSeriesAtDepth: 4 }],
       { behindText: false },
       undefined
     );

src/plugins/chart_expressions/expression_partition_vis/public/utils/layers/get_color.ts

@@ -21,12 +21,22 @@ const isTreemapOrMosaicChart = (shape: ChartTypes) =>
 
 export const byDataColorPaletteMap = (
   rows: Datatable['rows'],
-  columnId: string,
+  column: Partial<BucketColumns>,
   paletteDefinition: PaletteDefinition,
-  { params }: PaletteOutput
+  { params }: PaletteOutput,
+  formatters: Record<string, FieldFormat | undefined>,
+  formatter: FieldFormatsStart
 ) => {
   const colorMap = new Map<string, string | undefined>(
-    rows.map((item) => [String(item[columnId]), undefined])
+    rows.map((item) => {
+      const formattedName = getNodeLabel(
+        item[column.id ?? ''],
+        column,
+        formatters,
+        formatter.deserialize
+      );
+      return [formattedName, undefined];
+    })
   );
   let rankAtDepth = 0;
 

src/plugins/chart_expressions/expression_partition_vis/public/utils/layers/get_layers.ts

@@ -61,9 +61,11 @@ export const getLayers = (
   if (!syncColors && columns[1]?.id && palettes && visParams.palette) {
     byDataPalette = byDataColorPaletteMap(
       rows,
-      columns[1].id,
+      columns[1],
       palettes?.get(visParams.palette.name),
-      visParams.palette
+      visParams.palette,
+      formatters,
+      formatter
     );
   }