[Lens][Datatable] Fix color mapping of transposed datatables (#208623)

## Summary

In #189895 we added logic to the `getColorCategories` util function to
handle transpose columns. This mistakenly assumed the first row of the
datatable would include all transposed column ids (i.e.
`${value}---${columnId}`). After closer analysis this case is only
present the in datatable rendering (i.e. `table_basic.tsx`), but in this
context we also have the original non-transposed datatable.

So to simplify this we revert this logic to not care about transposed
datatables. Now the color mappings are correctly assigned across **split
by** columns.

<img width="720" alt="image"
src="https://github.com/user-attachments/assets/c588930e-53b9-409f-a257-2c5be35aaa38"
/>

Fixes #208555

### Checklist

- [x] [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
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

## Release note

Fixes an issue in Lens (#208555) Table where a split-by metric on a
terms rendered incorrect colors in table cells.

---------

Co-authored-by: Marta Bondyra <4283304+mbondyra@users.noreply.github.com>
Co-authored-by: Marco Liberati <dej611@users.noreply.github.com>

src/platform/packages/shared/chart-expressions-common/color_categories.test.ts

@@ -7,27 +7,38 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { DatatableRow } from '@kbn/expressions-plugin/common';
+import { DatatableColumn, DatatableRow } from '@kbn/expressions-plugin/common';
 import { getColorCategories } from './color_categories';
 
-const extensions = ['gz', 'css', '', 'rpm', 'deb', 'zip', null];
-const getExtension = (i: number) => extensions[i % extensions.length];
+const getNextExtension = (() => {
+  let i = 0;
+  const extensions = ['gz', 'css', '', 'rpm', 'deb', 'zip', null];
+  return () => extensions[i++ % extensions.length];
+})();
 
-const basicDatatableRows: DatatableRow[] = Array.from({ length: 30 }).map((_, i) => ({
-  count: i,
-  extension: getExtension(i),
-}));
-
-const isTransposedDatatableRows: DatatableRow[] = Array.from({ length: 30 }).map((_, i) => ({
-  count: i,
-  ['safari---extension']: getExtension(i),
-  ['chrome---extension']: getExtension(i + 1),
-  ['firefox---extension']: getExtension(i + 2),
-}));
+const basicDatatable = {
+  columns: ['count', 'extension'].map((id) => ({ id } as DatatableColumn)),
+  rows: Array.from({ length: 10 }).map((_, i) => ({
+    count: i,
+    extension: getNextExtension(),
+  })) as DatatableRow[],
+};
 
 describe('getColorCategories', () => {
-  it('should return all categories from datatable rows', () => {
-    expect(getColorCategories(basicDatatableRows, 'extension')).toEqual([
+  it('should return no categories when accessor is undefined', () => {
+    expect(getColorCategories(basicDatatable.rows)).toEqual([]);
+  });
+
+  it('should return no categories when accessor is not found', () => {
+    expect(getColorCategories(basicDatatable.rows, 'N/A')).toEqual([]);
+  });
+
+  it('should return no categories when no rows are defined', () => {
+    expect(getColorCategories(undefined, 'extension')).toEqual([]);
+  });
+
+  it('should return all categories from non-transpose datatable', () => {
+    expect(getColorCategories(basicDatatable.rows, 'extension')).toEqual([
       'gz',
       'css',
       '',
@@ -38,30 +49,8 @@ describe('getColorCategories', () => {
     ]);
   });
 
-  it('should exclude selected categories from datatable rows', () => {
-    expect(getColorCategories(basicDatatableRows, 'extension', false, ['', null])).toEqual([
-      'gz',
-      'css',
-      'rpm',
-      'deb',
-      'zip',
-    ]);
-  });
-
-  it('should return categories across all transpose columns of datatable rows', () => {
-    expect(getColorCategories(isTransposedDatatableRows, 'extension', true)).toEqual([
-      'gz',
-      'css',
-      '',
-      'rpm',
-      'deb',
-      'zip',
-      'null',
-    ]);
-  });
-
-  it('should exclude selected categories across all transpose columns of datatable rows', () => {
-    expect(getColorCategories(isTransposedDatatableRows, 'extension', true, ['', null])).toEqual([
+  it('should exclude selected categories from non-transpose datatable', () => {
+    expect(getColorCategories(basicDatatable.rows, 'extension', ['', null])).toEqual([
       'gz',
       'css',
       'rpm',

src/platform/packages/shared/chart-expressions-common/color_categories.ts

@@ -13,33 +13,27 @@ import { isMultiFieldKey } from '@kbn/data-plugin/common';
 /**
  * Get the stringified version of all the categories that needs to be colored in the chart.
  * Multifield keys will return as array of string and simple fields (numeric, string) will be returned as a plain unformatted string.
+ *
+ * Note: This does **NOT** support transposed columns
  */
 export function getColorCategories(
-  rows: DatatableRow[],
+  rows: DatatableRow[] = [],
   accessor?: string,
-  isTransposed?: boolean,
   exclude?: any[]
 ): Array<string | string[]> {
-  const ids = isTransposed
-    ? Object.keys(rows[0]).filter((key) => accessor && key.endsWith(accessor))
-    : accessor
-    ? [accessor]
-    : [];
+  if (!accessor) return [];
 
   return rows
-    .flatMap((r) =>
-      ids
-        .map((id) => r[id])
-        .filter((v) => !(v === undefined || exclude?.includes(v)))
-        .map((v) => {
-          // The categories needs to be stringified in their unformatted version.
-          // We can't distinguish between a number and a string from a text input and the match should
-          // work with both numeric field values and string values.
-          const key = (isMultiFieldKey(v) ? v.keys : [v]).map(String);
-          const stringifiedKeys = key.join(',');
-          return { key, stringifiedKeys };
-        })
-    )
+    .filter(({ [accessor]: v }) => !(v === undefined || exclude?.includes(v)))
+    .map((r) => {
+      const v = r[accessor];
+      // The categories needs to be stringified in their unformatted version.
+      // We can't distinguish between a number and a string from a text input and the match should
+      // work with both numeric field values and string values.
+      const key = (isMultiFieldKey(v) ? v.keys : [v]).map(String);
+      const stringifiedKeys = key.join(',');
+      return { key, stringifiedKeys };
+    })
     .reduce<{ keys: Set<string>; categories: Array<string | string[]> }>(
       (acc, { key, stringifiedKeys }) => {
         if (!acc.keys.has(stringifiedKeys)) {

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/components/dimension_editor.tsx

@@ -115,7 +115,7 @@ export function TableDimensionEditor(props: TableDimensionEditorProps) {
   };
   // need to tell the helper that the colorStops are required to display
   const displayStops = applyPaletteParams(props.paletteService, activePalette, currentMinMax);
-  const categories = getColorCategories(currentData?.rows ?? [], accessor, false, [null]);
+  const categories = getColorCategories(currentData?.rows, accessor, [null]);
 
   if (activePalette.name !== CUSTOM_PALETTE && activePalette.params?.stops) {
     activePalette.params.stops = applyPaletteParams(

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/components/table_basic.tsx

@@ -32,7 +32,7 @@ import { ClickTriggerEvent } from '@kbn/charts-plugin/public';
 import { IconChartDatatable } from '@kbn/chart-icons';
 import useObservable from 'react-use/lib/useObservable';
 import { getColorCategories } from '@kbn/chart-expressions-common';
-import { getOriginalId, isTransposeId } from '@kbn/transpose-utils';
+import { getOriginalId } from '@kbn/transpose-utils';
 import { CoreTheme } from '@kbn/core/public';
 import { getKbnPalettes } from '@kbn/palettes';
 import type { LensTableRowContextMenuEvent } from '../../../types';
@@ -405,16 +405,12 @@ export const DatatableComponent = (props: DatatableRenderProps) => {
       const dataType = getFieldMetaFromDatatable(firstLocalTable, originalId)?.type;
       const isBucketed = bucketedColumns.some((id) => id === columnId);
       const colorByTerms = shouldColorByTerms(dataType, isBucketed);
-
+      const categoryRows = (untransposedDataRef.current ?? firstLocalTable)?.rows;
       const data: ColorMappingInputData = colorByTerms
         ? {
             type: 'categories',
-            categories: getColorCategories(
-              firstLocalTable.rows,
-              originalId,
-              isTransposeId(columnId),
-              [null]
-            ),
+            // Must use non-transposed data here to correctly collate categories across transposed columns
+            categories: getColorCategories(categoryRows, originalId, [null]),
           }
         : {
             type: 'ranges',

x-pack/platform/plugins/shared/lens/public/visualizations/partition/dimension_editor.tsx

@@ -130,7 +130,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
     currentLayer.colorMapping
   );
   const table = props.frame.activeData?.[currentLayer.layerId];
-  const splitCategories = getColorCategories(table?.rows ?? [], props.accessor);
+  const splitCategories = getColorCategories(table?.rows, props.accessor);
 
   return (
     <>

x-pack/platform/plugins/shared/lens/public/visualizations/tagcloud/tags_dimension_editor.tsx

@@ -65,7 +65,7 @@ export function TagsDimensionEditor({
     state.colorMapping
   );
   const table = frame.activeData?.[state.layerId];
-  const splitCategories = getColorCategories(table?.rows ?? [], state.tagAccessor);
+  const splitCategories = getColorCategories(table?.rows, state.tagAccessor);
 
   const setColorMapping = useCallback(
     (colorMapping?: ColorMapping.Config) => {

x-pack/platform/plugins/shared/lens/public/visualizations/xy/xy_config_panel/dimension_editor.tsx

@@ -134,7 +134,7 @@ export function DataDimensionEditor(
 
   const table = props.frame.activeData?.[layer.layerId];
   const { splitAccessor } = layer;
-  const splitCategories = getColorCategories(table?.rows ?? [], splitAccessor);
+  const splitCategories = getColorCategories(table?.rows, splitAccessor);
 
   if (props.groupId === 'breakdown') {
     return !layer.collapseFn ? (