[ColorMapping] Store assignments as raw/serialized values (#207957)

## Summary

Refactors color mapping logic to store raw values in save objects.

Main changes:

- Instead assignments having a single rule with many values, we now have
assignments with multiple rules of varying types with a single value.
- As per the previous change, there is no more explicit `auto` rule,
this is now implicit when there are no rules in the assignment.
- Raw values can now be raw and in cases such as `RangeKey`s and
`MultiValueKey`s can be instance values. This is not ok for storing in
redux nor the SO so be have added a `serialize` methods to facilitate
this value storage. Code changes attempt to identify which value is used
based on variable/param names and type aliases as both are `unknown`.
- For values that are non-string or non-number types, we no longer allow
creating custom matching options.
- Values are now correctly formatted.
- Add runtime migration for `xy`, `partition`, `tagcloud` and
`datatable` viz. Requires `formBased` `datasourceState` to determine
best string-to-raw value convertion.

Closes #193080
Fixes #187519

### Checklist

- [x] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [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] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [x] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [x] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [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)

### Identify risks

<!-- Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging -->

- [x] Migration from previous stored types.

## Release note

This PR fixes an issue
([#193080](https://github.com/elastic/kibana/issues/193080)) where
custom ranges and multi-field values were not correctly colored based on
selected color mapping configurations.

This change includes a runtime migration to convert old mappings as
strings to their raw value equivalent. This conversion is done at
runtime when the vis is rendered and only updated when the visualization
is saved. Thus this conversion does not dirty the state of the
visualization such as when first opening to edit.

This _should_ have no affect to the user apart from improved value
formatting in the color mapping assignment selection UI. In rare cases,
some assignments may not be correctly converted exactly to the new raw
value but are still preserved to use as string value matches. The only
know case where we are not be able to fully convert the value is when
using labels on custom ranges, the label will not show in the color
mapping assignment UI unless the value is removed as added back.

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Co-authored-by: Marco Liberati <dej611@users.noreply.github.com>
Co-authored-by: Marco Vettorello <marco.vettorello@elastic.co>

packages/kbn-optimizer/limits.yml

@@ -10,7 +10,7 @@ pageLoadAssetSize:
   banners: 17946
   canvas: 29355
   cases: 180037
-  charts: 55000
+  charts: 60000
   cloud: 21076
   cloudDataMigration: 19170
   cloudExperiments: 109746

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

@@ -7,55 +7,148 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { DatatableColumn, DatatableRow } from '@kbn/expressions-plugin/common';
-import { getColorCategories } from './color_categories';
+import { DatatableRow } from '@kbn/expressions-plugin/common';
+import { getColorCategories, getLegacyColorCategories } from './color_categories';
+import { MultiFieldKey, RangeKey } from '@kbn/data-plugin/common';
 
-const getNextExtension = (() => {
-  let i = 0;
-  const extensions = ['gz', 'css', '', 'rpm', 'deb', 'zip', null];
-  return () => extensions[i++ % extensions.length];
-})();
+class FakeClass {}
 
-const basicDatatable = {
-  columns: ['count', 'extension'].map((id) => ({ id } as DatatableColumn)),
-  rows: Array.from({ length: 10 }).map((_, i) => ({
-    count: i,
-    extension: getNextExtension(),
-  })) as DatatableRow[],
-};
+const values = [
+  1,
+  false,
+  true,
+  0,
+  NaN,
+  null,
+  undefined,
+  '',
+  'test-string',
+  { test: 'obj' },
+  ['array'],
+];
+const mockRange = new RangeKey({ from: 0, to: 100 });
+const mockMultiField = new MultiFieldKey({ key: ['one', 'two'] });
+const fakeClass = new FakeClass();
+const complex = [mockRange, mockMultiField, fakeClass];
 
-describe('getColorCategories', () => {
-  it('should return no categories when accessor is undefined', () => {
-    expect(getColorCategories(basicDatatable.rows)).toEqual([]);
+const mockDatatableRows = Array.from({ length: 20 }).map<DatatableRow>((_, i) => ({
+  count: i,
+  value: values[i % values.length],
+  complex: complex[i % complex.length],
+}));
+
+describe('Color Categories', () => {
+  describe('getColorCategories', () => {
+    it('should return no categories when accessor is undefined', () => {
+      expect(getColorCategories(mockDatatableRows)).toEqual([]);
+    });
+
+    it('should return no categories when accessor is not found', () => {
+      expect(getColorCategories(mockDatatableRows, 'N/A')).toEqual([]);
+    });
+
+    it('should return no categories when no rows are defined', () => {
+      expect(getColorCategories(undefined, 'extension')).toEqual([]);
+    });
+
+    it('should return all categories from mixed value datatable', () => {
+      expect(getColorCategories(mockDatatableRows, 'value')).toEqual([
+        1,
+        false,
+        true,
+        0,
+        NaN,
+        null,
+        undefined,
+        '',
+        'test-string',
+        {
+          test: 'obj',
+        },
+        ['array'],
+      ]);
+    });
+
+    it('should exclude selected categories from datatable', () => {
+      expect(
+        getColorCategories(mockDatatableRows, 'value', [
+          1,
+          false,
+          true,
+          0,
+          NaN,
+          null,
+          undefined,
+          '',
+        ])
+      ).toEqual([
+        'test-string',
+        {
+          test: 'obj',
+        },
+        ['array'],
+      ]);
+    });
+
+    it('should return known serialized categories from datatable', () => {
+      expect(getColorCategories(mockDatatableRows, 'complex', [])).toEqual([
+        mockRange.serialize(),
+        mockMultiField.serialize(),
+        fakeClass,
+      ]);
+    });
   });
 
-  it('should return no categories when accessor is not found', () => {
-    expect(getColorCategories(basicDatatable.rows, 'N/A')).toEqual([]);
-  });
+  describe('getLegacyColorCategories', () => {
+    it('should return no categories when accessor is undefined', () => {
+      expect(getLegacyColorCategories(mockDatatableRows)).toEqual([]);
+    });
 
-  it('should return no categories when no rows are defined', () => {
-    expect(getColorCategories(undefined, 'extension')).toEqual([]);
-  });
+    it('should return no categories when accessor is not found', () => {
+      expect(getLegacyColorCategories(mockDatatableRows, 'N/A')).toEqual([]);
+    });
 
-  it('should return all categories from non-transpose datatable', () => {
-    expect(getColorCategories(basicDatatable.rows, 'extension')).toEqual([
-      'gz',
-      'css',
-      '',
-      'rpm',
-      'deb',
-      'zip',
-      'null',
-    ]);
-  });
+    it('should return no categories when no rows are defined', () => {
+      expect(getLegacyColorCategories(undefined, 'extension')).toEqual([]);
+    });
 
-  it('should exclude selected categories from non-transpose datatable', () => {
-    expect(getColorCategories(basicDatatable.rows, 'extension', ['', null])).toEqual([
-      'gz',
-      'css',
-      'rpm',
-      'deb',
-      'zip',
-    ]);
+    it('should return all categories from mixed value datatable', () => {
+      expect(getLegacyColorCategories(mockDatatableRows, 'value')).toEqual([
+        '1',
+        'false',
+        'true',
+        '0',
+        'NaN',
+        'null',
+        'undefined',
+        '',
+        'test-string',
+        '{"test":"obj"}',
+        'array',
+      ]);
+    });
+
+    it('should exclude selected categories from datatable', () => {
+      expect(
+        getLegacyColorCategories(mockDatatableRows, 'value', [
+          1,
+          false,
+          true,
+          0,
+          NaN,
+          null,
+          undefined,
+          '',
+        ])
+      ).toEqual(['test-string', '{"test":"obj"}', 'array']);
+    });
+
+    it('should return known serialized categories from datatable', () => {
+      expect(getLegacyColorCategories(mockDatatableRows, 'complex', [])).toEqual([
+        String(mockRange),
+        String(mockMultiField),
+        JSON.stringify(fakeClass),
+      ]);
+    });
   });
 });

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

@@ -8,40 +8,46 @@
  */
 
 import { DatatableRow } from '@kbn/expressions-plugin/common';
-import { isMultiFieldKey } from '@kbn/data-plugin/common';
+import { RawValue, SerializedValue, serializeField } from '@kbn/data-plugin/common';
+import { getValueKey } from '@kbn/coloring';
 
 /**
- * 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.
+ * Returns all serialized categories of the dataset for color matching.
+ * All non-serializable fields will be as a plain unformatted string.
  *
  * Note: This does **NOT** support transposed columns
  */
 export function getColorCategories(
   rows: DatatableRow[] = [],
   accessor?: string,
-  exclude?: any[]
-): Array<string | string[]> {
+  exclude?: RawValue[],
+  legacyMode: boolean = false // stringifies raw values
+): SerializedValue[] {
   if (!accessor) return [];
 
-  return rows
-    .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)) {
-          acc.keys.add(stringifiedKeys);
-          acc.categories.push(key.length === 1 ? key[0] : key);
-        }
-        return acc;
-      },
-      { keys: new Set(), categories: [] }
-    ).categories;
+  const seen = new Set<unknown>();
+  return rows.reduce<SerializedValue[]>((acc, row) => {
+    const hasValue = Object.hasOwn(row, accessor);
+    const rawValue: RawValue = row[accessor];
+    const key = getValueKey(rawValue);
+    if (hasValue && !exclude?.includes(rawValue) && !seen.has(key)) {
+      const value = serializeField(rawValue);
+      seen.add(key);
+      acc.push(legacyMode ? key : value);
+    }
+    return acc;
+  }, []);
+}
+
+/**
+ * Returns all *stringified* categories of the dataset for color matching.
+ *
+ * Should **only** be used with legacy `palettes`
+ */
+export function getLegacyColorCategories(
+  rows?: DatatableRow[],
+  accessor?: string,
+  exclude?: RawValue[]
+): string[] {
+  return getColorCategories(rows, accessor, exclude, true).map(String);
 }

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

@@ -15,4 +15,4 @@ export {
 } from './utils';
 export type { Simplify, MakeOverridesSerializable, ChartSizeSpec, ChartSizeEvent } from './types';
 export { isChartSizeEvent } from './types';
-export { getColorCategories } from './color_categories';
+export { getColorCategories, getLegacyColorCategories } from './color_categories';

src/platform/packages/shared/chart-expressions-common/tsconfig.json

@@ -19,5 +19,6 @@
     "@kbn/core-execution-context-common",
     "@kbn/expressions-plugin",
     "@kbn/data-plugin",
+    "@kbn/coloring",
   ]
 }

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/README.md

@@ -3,9 +3,9 @@
 This shared component can be used to define a color mapping as an association of one or multiple string values to a color definition.
 
 This package provides:
-- a React component, called `CategoricalColorMapping` that provides a simplified UI (that in general can be hosted in a flyout), that helps the user generate a `ColorMapping.Config` object that descibes the mappings configuration
-- a function `getColorFactory` that given a color mapping configuration returns a function that maps a passed category to the corresponding color
-- a definition scheme for the color mapping, based on the type `ColorMapping.Config`, that provides an extensible way of describing the link between colors and rules. Collects the minimal information required apply colors based on categories. Together with the `ColorMappingInputData` can be used to get colors in a deterministic way.
+- A React component, called `CategoricalColorMapping` that provides a simplified UI (that in general can be hosted in a flyout), that helps the user generate a `ColorMapping.Config` object that describes the mappings configuration
+- A function `getColorFactory` that given a color mapping configuration returns a function that maps a passed category to the corresponding color
+- A definition scheme for the color mapping, based on the type `ColorMapping.Config`, that provides an extensible way of describing the link between colors and rules. Collects the minimal information required apply colors based on categories. Together with the `ColorMappingInputData` can be used to get colors in a deterministic way.
 
 
 An example of the configuration is the following:
@@ -14,10 +14,10 @@ const DEFAULT_COLOR_MAPPING_CONFIG: ColorMapping.Config = {
   assignmentMode: 'auto',
   assignments: [
     {
-        rule: {
-            type: 'matchExactly',
-            values: [''];
-        },
+        rules: [{
+            type: 'match',
+            pattern: '';
+        }],
         color: {
             type: 'categorical',
             paletteId: 'eui',
@@ -27,9 +27,9 @@ const DEFAULT_COLOR_MAPPING_CONFIG: ColorMapping.Config = {
   ],
   specialAssignments: [
     {
-      rule: {
+      rules: [{
         type: 'other',
-      },
+      }],
       color: {
         type: 'categorical',
         paletteId: 'neutral',
@@ -45,7 +45,7 @@ const DEFAULT_COLOR_MAPPING_CONFIG: ColorMapping.Config = {
 };
 ```
 
-The function `getColorFactory` is a curry function where, given the model, a palette getter, the theme mode (dark/light) and a list of categories, returns a function that can be used to pick the right color based on a given category.
+The function `getColorFactory` is a curry function where, given the model, a palette getter, the theme mode (dark/light) and a list of categories, returns a `ColorHandlingFn` that can be used to pick the right color based on a given category.
 
 ```ts
 function getColorFactory(
@@ -56,22 +56,19 @@ function getColorFactory(
       type: 'categories';
       categories: Array<string | string[]>;
     }
-): (category: string | string[]) => Color
+): ColorHandlingFn
 ```
 
-
-
 A `category` can be in the shape of a plain string or an array of strings. Numbers, MultiFieldKey, IP etc needs to be stringified.
 
-
 The `CategoricalColorMapping` React component has the following props:
 
 ```tsx
 function CategoricalColorMapping(props: {
   /** The initial color mapping model, usually coming from a the visualization saved object */
   model: ColorMapping.Config;
-  /** A map of paletteId and palette configuration */
-  palettes: Map<string, ColorMapping.CategoricalPalette>;
+  /** A collection of palette configurations */
+  palettes: KbnPalettes;
   /** A data description of what needs to be colored */
   data: ColorMappingInputData;
   /** Theme dark mode */
@@ -80,8 +77,11 @@ function CategoricalColorMapping(props: {
   specialTokens: Map<string, string>;
   /** A function called at every change in the model */
   onModelUpdate: (model: ColorMapping.Config) => void;
+  /** Formatter for raw value assignments */
+  formatter?: IFieldFormat;
+  /** Allow custom match rule when no other option is found */
+  allowCustomMatch?: boolean;
 })
-
 ```
 
-the `onModelUpdate` callback is called everytime a change in the model is applied from within the component. Is not called when the `model` prop is updated.
+the `onModelUpdate` callback is called every time a change in the model is applied from within the component. Is not called when the `model` prop is updated.

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/__stories__/color_mapping.stories.tsx

@@ -12,11 +12,11 @@ import { getKbnPalettes } from '@kbn/palettes';
 import { EuiFlyout, EuiForm, EuiPage, isColorDark } from '@elastic/eui';
 import type { StoryFn } from '@storybook/react';
 import { css } from '@emotion/react';
+import { RawValue, deserializeField } from '@kbn/data-plugin/common';
 import { CategoricalColorMapping, ColorMappingProps } from '../categorical_color_mapping';
 import { DEFAULT_COLOR_MAPPING_CONFIG } from '../config/default_color_mapping';
 import { ColorMapping } from '../config';
 import { getColorFactory } from '../color/color_handling';
-import { ruleMatch } from '../color/rule_matching';
 import { getValidColor } from '../color/color_math';
 
 export default {
@@ -25,6 +25,8 @@ export default {
   decorators: [(story: Function) => story()],
 };
 
+const formatter = (value: unknown) => String(value);
+
 const Template: StoryFn<FC<ColorMappingProps>> = (args) => {
   const [updatedModel, setUpdateModel] = useState<ColorMapping.Config>(
     DEFAULT_COLOR_MAPPING_CONFIG
@@ -37,11 +39,12 @@ const Template: StoryFn<FC<ColorMappingProps>> = (args) => {
     <EuiPage>
       <ol>
         {args.data.type === 'categories' &&
-          args.data.categories.map((c, i) => {
-            const match = updatedModel.assignments.some(({ rule }) => {
-              return ruleMatch(rule, c);
-            });
-            const color = colorFactory(c);
+          args.data.categories.map((category, i) => {
+            const value: RawValue = deserializeField(category);
+            const match = updatedModel.assignments.some(({ rules }) =>
+              rules.some((r) => (r.type === 'raw' ? r.value === value : false))
+            );
+            const color = colorFactory(value);
             const isDark = isColorDark(...getValidColor(color).rgb());
 
             return (
@@ -58,7 +61,7 @@ const Template: StoryFn<FC<ColorMappingProps>> = (args) => {
                   font-weight: ${match ? 'bold' : 'normal'};
                 `}
               >
-                {c}
+                {formatter(value)}
               </li>
             );
           })}
@@ -90,9 +93,11 @@ export const Default = {
       },
       specialAssignments: [
         {
-          rule: {
-            type: 'other',
-          },
+          rules: [
+            {
+              type: 'other',
+            },
+          ],
           color: {
             type: 'loop',
           },

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/__stories__/raw_color_mapping.stories.tsx

@@ -0,0 +1,141 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import React, { FC, useState } from 'react';
+import { getKbnPalettes } from '@kbn/palettes';
+import { EuiFlyout, EuiForm, EuiPage, isColorDark } from '@elastic/eui';
+import type { StoryFn } from '@storybook/react';
+import { css } from '@emotion/react';
+import {
+  MultiFieldKey,
+  RawValue,
+  SerializedValue,
+  deserializeField,
+} from '@kbn/data-plugin/common';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
+import { CategoricalColorMapping, ColorMappingProps } from '../categorical_color_mapping';
+import { DEFAULT_COLOR_MAPPING_CONFIG } from '../config/default_color_mapping';
+import { ColorMapping } from '../config';
+import { getColorFactory } from '../color/color_handling';
+import { getValidColor } from '../color/color_math';
+
+export default {
+  title: 'Raw Color Mapping',
+  component: CategoricalColorMapping,
+  decorators: [(story: Function) => story()],
+};
+
+const formatter = {
+  convert: (value: MultiFieldKey) => {
+    return value.keys.join(' - ');
+  },
+} as IFieldFormat;
+
+const Template: StoryFn<FC<ColorMappingProps>> = (args) => {
+  const [updatedModel, setUpdateModel] = useState<ColorMapping.Config>(
+    DEFAULT_COLOR_MAPPING_CONFIG
+  );
+
+  const palettes = getKbnPalettes({ name: 'amsterdam', darkMode: false });
+  const colorFactory = getColorFactory(updatedModel, palettes, false, args.data);
+
+  return (
+    <EuiPage>
+      <ol>
+        {args.data.type === 'categories' &&
+          args.data.categories.map((category: SerializedValue, i) => {
+            const value: RawValue = deserializeField(category);
+            const match = updatedModel.assignments.some(({ rules }) =>
+              rules.some((r) => (r.type === 'raw' ? String(r.value) === String(category) : false))
+            );
+            const color = colorFactory(value);
+            const isDark = isColorDark(...getValidColor(color).rgb());
+
+            return (
+              <li
+                key={i}
+                css={css`
+                  width: ${100 + 200 * Math.abs(Math.cos(i))}px;
+                  height: 30px;
+                  margin: 2px;
+                  padding: 5px;
+                  background: ${color};
+                  color: ${isDark ? 'white' : 'black'};
+                  border: ${match ? '2px solid black' : 'none'};
+                  font-weight: ${match ? 'bold' : 'normal'};
+                `}
+              >
+                {formatter.convert(value)}
+              </li>
+            );
+          })}
+      </ol>
+      <EuiFlyout
+        css={{ width: 350, minInlineSize: 366, padding: '8px', overflow: 'auto' }}
+        onClose={() => {}}
+        hideCloseButton
+        ownFocus={false}
+      >
+        <EuiForm>
+          <CategoricalColorMapping {...args} palettes={palettes} onModelUpdate={setUpdateModel} />
+        </EuiForm>
+      </EuiFlyout>
+    </EuiPage>
+  );
+};
+
+export const Default = {
+  render: Template,
+  args: {
+    model: {
+      ...DEFAULT_COLOR_MAPPING_CONFIG,
+      paletteId: 'eui_amsterdam',
+
+      colorMode: {
+        type: 'categorical',
+      },
+      specialAssignments: [
+        {
+          rules: [
+            {
+              type: 'other',
+            },
+          ],
+          color: {
+            type: 'loop',
+          },
+          touched: false,
+        },
+      ],
+      assignments: [],
+    },
+    isDarkMode: false,
+    formatter,
+    data: {
+      type: 'categories',
+      categories: [
+        { type: 'multiFieldKey', keys: ['US', 'Canada'] },
+        { type: 'multiFieldKey', keys: ['Mexico'] },
+        { type: 'multiFieldKey', keys: ['Brasil'] },
+        { type: 'multiFieldKey', keys: ['Canada'] },
+        { type: 'multiFieldKey', keys: ['Canada', 'US'] },
+        { type: 'multiFieldKey', keys: ['Italy', 'Germany'] },
+        { type: 'multiFieldKey', keys: ['France'] },
+        { type: 'multiFieldKey', keys: ['Spain', 'Portugal'] },
+        { type: 'multiFieldKey', keys: ['UK'] },
+        { type: 'multiFieldKey', keys: ['Sweden'] },
+        { type: 'multiFieldKey', keys: ['Sweden', 'Finland'] },
+      ],
+    },
+
+    specialTokens: new Map(),
+    // eslint-disable-next-line no-console
+    onModelUpdate: (model: any) => console.log(model),
+  },
+};

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/categorical_color_mapping.test.tsx

@@ -8,105 +8,113 @@
  */
 
 import React from 'react';
-import { mount } from 'enzyme';
-import { CategoricalColorMapping, ColorMappingInputData } from './categorical_color_mapping';
-import { DEFAULT_COLOR_MAPPING_CONFIG } from './config/default_color_mapping';
-import { MULTI_FIELD_KEY_SEPARATOR } from '@kbn/data-plugin/common';
+import { fireEvent, render, screen } from '@testing-library/react';
+import '@testing-library/jest-dom';
+
+import { fieldFormatsServiceMock } from '@kbn/field-formats-plugin/public/mocks';
 import { getKbnPalettes } from '@kbn/palettes';
 
-const ASSIGNMENTS_LIST = '[data-test-subj="lns-colorMapping-assignmentsList"]';
-const ASSIGNMENTS_PROMPT = '[data-test-subj="lns-colorMapping-assignmentsPrompt"]';
-const ASSIGNMENTS_PROMPT_ADD_ALL = '[data-test-subj="lns-colorMapping-assignmentsPromptAddAll"]';
-const ASSIGNMENT_ITEM = (i: number) => `[data-test-subj="lns-colorMapping-assignmentsItem${i}"]`;
+import {
+  CategoricalColorMapping,
+  ColorMappingInputCategoricalData,
+  ColorMappingInputData,
+  ColorMappingProps,
+} from './categorical_color_mapping';
+import { DEFAULT_COLOR_MAPPING_CONFIG } from './config/default_color_mapping';
+
+const ASSIGNMENTS_LIST = 'lns-colorMapping-assignmentsList';
+const ASSIGNMENTS_PROMPT = 'lns-colorMapping-assignmentsPrompt';
+const ASSIGNMENTS_PROMPT_ADD_ALL = 'lns-colorMapping-assignmentsPromptAddAll';
+const ASSIGNMENT_ITEM = (i: number) => `lns-colorMapping-assignmentsItem${i}`;
+
+const palettes = getKbnPalettes({ name: 'amsterdam', darkMode: false });
+const specialTokens = new Map([
+  ['__other__', 'Other'],
+  ['__empty__', '(Empty)'],
+  ['', '(Empty)'],
+]);
+const categoryData: ColorMappingInputCategoricalData = {
+  type: 'categories',
+  categories: ['categoryA', 'categoryB'],
+};
+const mockFormatter = fieldFormatsServiceMock.createStartContract().deserialize();
 
 describe('color mapping', () => {
-  const palettes = getKbnPalettes({ name: 'amsterdam', darkMode: false });
+  let defaultProps: ColorMappingProps;
+
+  mockFormatter.convert = jest.fn(
+    (v: any) => (typeof v === 'string' ? specialTokens.get(v) ?? v : JSON.stringify(v)) // simple way to check formatting is applied
+  );
+  const onModelUpdateFn = jest.fn();
+
+  beforeEach(() => {
+    defaultProps = {
+      data: categoryData,
+      isDarkMode: false,
+      model: { ...DEFAULT_COLOR_MAPPING_CONFIG },
+      palettes,
+      onModelUpdate: onModelUpdateFn,
+      specialTokens,
+      formatter: mockFormatter,
+    };
+  });
+
+  const renderCategoricalColorMapping = (props: Partial<ColorMappingProps> = {}) => {
+    return render(<CategoricalColorMapping {...defaultProps} {...props} />);
+  };
 
   it('load a default color mapping', () => {
-    const dataInput: ColorMappingInputData = {
-      type: 'categories',
-      categories: ['categoryA', 'categoryB'],
-    };
-    const onModelUpdateFn = jest.fn();
-    const component = mount(
-      <CategoricalColorMapping
-        data={dataInput}
-        isDarkMode={false}
-        model={{ ...DEFAULT_COLOR_MAPPING_CONFIG }}
-        palettes={palettes}
-        onModelUpdate={onModelUpdateFn}
-        specialTokens={new Map()}
-      />
-    );
+    renderCategoricalColorMapping();
 
     // empty list prompt visible
-    expect(component.find(ASSIGNMENTS_PROMPT)).toBeTruthy();
-    expect(onModelUpdateFn).not.toBeCalled();
+    expect(screen.getByTestId(ASSIGNMENTS_PROMPT)).toBeInTheDocument();
+    expect(onModelUpdateFn).not.toHaveBeenCalled();
   });
 
   it('Add all terms to assignments', () => {
-    const dataInput: ColorMappingInputData = {
-      type: 'categories',
-      categories: ['categoryA', 'categoryB'],
-    };
-    const onModelUpdateFn = jest.fn();
-    const component = mount(
-      <CategoricalColorMapping
-        data={dataInput}
-        isDarkMode={false}
-        model={{ ...DEFAULT_COLOR_MAPPING_CONFIG }}
-        palettes={palettes}
-        onModelUpdate={onModelUpdateFn}
-        specialTokens={new Map()}
-      />
-    );
-    component.find(ASSIGNMENTS_PROMPT_ADD_ALL).hostNodes().simulate('click');
-    expect(onModelUpdateFn).toBeCalledTimes(1);
-    expect(component.find(ASSIGNMENTS_LIST).hostNodes().children().length).toEqual(
-      dataInput.categories.length
-    );
-    dataInput.categories.forEach((category, index) => {
-      const assignment = component.find(ASSIGNMENT_ITEM(index)).hostNodes();
-      expect(assignment.text()).toEqual(category);
-      expect(assignment.hasClass('euiComboBox-isDisabled')).toEqual(false);
+    renderCategoricalColorMapping();
+
+    fireEvent.click(screen.getByTestId(ASSIGNMENTS_PROMPT_ADD_ALL));
+
+    expect(onModelUpdateFn).toHaveBeenCalledTimes(1);
+    const assignmentsList = screen.getByTestId(ASSIGNMENTS_LIST);
+    expect(assignmentsList.children.length).toEqual(categoryData.categories.length);
+
+    categoryData.categories.forEach((category, index) => {
+      const assignment = screen.getByTestId(ASSIGNMENT_ITEM(index));
+      expect(assignment).toHaveTextContent(String(category));
+      expect(assignment).not.toHaveClass('euiComboBox-isDisabled');
     });
   });
 
   it('handle special tokens, multi-fields keys and non-trimmed whitespaces', () => {
-    const dataInput: ColorMappingInputData = {
+    const data: ColorMappingInputData = {
       type: 'categories',
-      categories: ['__other__', ['fieldA', 'fieldB'], '__empty__', '   with-whitespaces   '],
+      categories: [
+        '__other__',
+        '__empty__',
+        '',
+        '   with-whitespaces   ',
+        { type: 'multiFieldKey', keys: ['gz', 'CN'] },
+        { type: 'rangeKey', from: 0, to: 1000, ranges: [{ from: 0, to: 1000, label: '' }] },
+      ],
     };
-    const onModelUpdateFn = jest.fn();
-    const component = mount(
-      <CategoricalColorMapping
-        data={dataInput}
-        isDarkMode={false}
-        model={{ ...DEFAULT_COLOR_MAPPING_CONFIG }}
-        palettes={palettes}
-        onModelUpdate={onModelUpdateFn}
-        specialTokens={
-          new Map([
-            ['__other__', 'Other'],
-            ['__empty__', '(Empty)'],
-          ])
-        }
-      />
+    renderCategoricalColorMapping({ data });
+
+    fireEvent.click(screen.getByTestId(ASSIGNMENTS_PROMPT_ADD_ALL));
+
+    const assignmentsList = screen.getByTestId(ASSIGNMENTS_LIST);
+    expect(assignmentsList.children.length).toEqual(data.categories.length);
+
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(0))).toHaveTextContent('Other');
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(1))).toHaveTextContent('(Empty)');
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(2))).toHaveTextContent('(Empty)');
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(3))).toHaveTextContent('   with-whitespaces   ', {
+      normalizeWhitespace: false,
+    });
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(4))).toHaveTextContent('{"keys":["gz","CN"]}');
+    expect(screen.getByTestId(ASSIGNMENT_ITEM(5))).toHaveTextContent(
+      '{"gte":0,"lt":1000,"label":""}'
     );
-    component.find(ASSIGNMENTS_PROMPT_ADD_ALL).hostNodes().simulate('click');
-    expect(component.find(ASSIGNMENTS_LIST).hostNodes().children().length).toEqual(
-      dataInput.categories.length
-    );
-    const assignment1 = component.find(ASSIGNMENT_ITEM(0)).hostNodes();
-    expect(assignment1.text()).toEqual('Other');
-
-    const assignment2 = component.find(ASSIGNMENT_ITEM(1)).hostNodes();
-    expect(assignment2.text()).toEqual(`fieldA${MULTI_FIELD_KEY_SEPARATOR}fieldB`);
-
-    const assignment3 = component.find(ASSIGNMENT_ITEM(2)).hostNodes();
-    expect(assignment3.text()).toEqual('(Empty)');
-
-    const assignment4 = component.find(ASSIGNMENT_ITEM(3)).hostNodes();
-    expect(assignment4.text()).toEqual('   with-whitespaces   ');
   });
 });

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/categorical_color_mapping.tsx

@@ -12,6 +12,8 @@ import { Provider } from 'react-redux';
 import { type EnhancedStore, configureStore } from '@reduxjs/toolkit';
 import { isEqual } from 'lodash';
 import { KbnPalettes } from '@kbn/palettes';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
+import { SerializedValue } from '@kbn/data-plugin/common';
 import { colorMappingReducer, updateModel } from './state/color_mapping';
 import { Container } from './components/container/container';
 import { ColorMapping } from './config';
@@ -19,8 +21,10 @@ import { uiReducer } from './state/ui';
 
 export interface ColorMappingInputCategoricalData {
   type: 'categories';
-  /** an ORDERED array of categories rendered in the visualization  */
-  categories: Array<string | string[]>;
+  /**
+   * An **ordered** array of serialized categories rendered in the visualization
+   */
+  categories: SerializedValue[];
 }
 
 export interface ColorMappingInputContinuousData {
@@ -42,18 +46,38 @@ export type ColorMappingInputData =
  * The props of the CategoricalColorMapping component
  */
 export interface ColorMappingProps {
-  /** The initial color mapping model, usually coming from a the visualization saved object */
+  /**
+   * The initial color mapping model, usually coming from a the visualization saved object
+   */
   model: ColorMapping.Config;
-  /** A collection of palette configurations */
+  /**
+   * A collection of palette configurations
+   */
   palettes: KbnPalettes;
-  /** A data description of what needs to be colored */
+  /**
+   * A data description of what needs to be colored
+   */
   data: ColorMappingInputData;
-  /** Theme dark mode */
+  /**
+   * Theme dark mode
+   */
   isDarkMode: boolean;
-  /** A map between original and formatted tokens used to handle special cases, like the Other bucket and the empty bucket */
+  /**
+   * A map between original and formatted tokens used to handle special cases, like the Other bucket and the empty bucket
+   */
   specialTokens: Map<string, string>;
-  /** A function called at every change in the model */
+  /**
+   * A function called at every change in the model
+   */
   onModelUpdate: (model: ColorMapping.Config) => void;
+  /**
+   * Formatter for raw value assignments
+   */
+  formatter?: IFieldFormat;
+  /**
+   * Allow custom match rule when no other option is found
+   */
+  allowCustomMatch?: boolean;
 }
 
 /**
@@ -88,7 +112,7 @@ export class CategoricalColorMapping extends React.Component<ColorMappingProps>
     }
   }
   render() {
-    const { palettes, data, isDarkMode, specialTokens } = this.props;
+    const { palettes, data, isDarkMode, specialTokens, formatter, allowCustomMatch } = this.props;
     return (
       <Provider store={this.store}>
         <Container
@@ -96,6 +120,8 @@ export class CategoricalColorMapping extends React.Component<ColorMappingProps>
           data={data}
           isDarkMode={isDarkMode}
           specialTokens={specialTokens}
+          formatter={formatter}
+          allowCustomMatch={allowCustomMatch}
         />
       </Provider>
     );

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/color_assignment_matcher.ts

@@ -0,0 +1,106 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { RawValue, deserializeField } from '@kbn/data-plugin/common';
+import { ColorMapping } from '../config';
+import { getValueKey } from './utils';
+
+type AssignmentMatchCount = [assignmentIndex: number, matchCount: number];
+
+/**
+ * A class to encapsulate assignment logic
+ */
+export class ColorAssignmentMatcher {
+  /**
+   * Reference to original assignments
+   */
+  readonly #assignments: ColorMapping.Assignment[];
+
+  /**
+   * Map values (or keys) to assignment index and match count
+   */
+  #assignmentMap: Map<string, AssignmentMatchCount>;
+
+  constructor(assignments: ColorMapping.Assignment[]) {
+    this.#assignments = assignments;
+    this.#assignmentMap = this.#assignments.reduce<Map<string, AssignmentMatchCount>>(
+      (acc, assignment, i) => {
+        assignment.rules.forEach((rule) => {
+          const key = getKey(rule);
+          if (key !== null) {
+            const [index = i, matchCount = 0] = acc.get(key) ?? [];
+            acc.set(key, [index, matchCount + 1]);
+          }
+        });
+        return acc;
+      },
+      new Map()
+    );
+  }
+
+  #getMatch(value: RawValue): AssignmentMatchCount {
+    const key = getValueKey(value);
+    return this.#assignmentMap.get(key) ?? [-1, 0];
+  }
+
+  /**
+   * Returns count of matching assignments for given value
+   */
+  getCount(value: RawValue) {
+    const [, count] = this.#getMatch(value);
+    return count;
+  }
+
+  /**
+   * Returns true if given value has multiple matching assignment
+   */
+  hasDuplicate(value: RawValue) {
+    const [, count] = this.#getMatch(value);
+    return count > 1;
+  }
+
+  /**
+   * Returns true if given value has matching assignment
+   */
+  hasMatch(value: RawValue) {
+    return this.getCount(value) > 0;
+  }
+
+  /**
+   * Returns index of first matching assignment for given value
+   */
+  getIndex(value: RawValue) {
+    const [index] = this.#getMatch(value);
+    return index;
+  }
+}
+
+function getKey(rule: ColorMapping.ColorRule): string | null {
+  if (rule.type === 'match' && rule.matchEntireWord) {
+    return rule.matchCase ? rule.pattern : rule.pattern.toLowerCase();
+  }
+
+  if (rule.type === 'raw') {
+    return getValueKey(deserializeField(rule.value));
+  }
+
+  // nondeterministic match, cannot assign ambiguous keys
+  // requires pattern matching all previous rules
+  return null;
+}
+
+/**
+ * A simplified map to track assignment match counts
+ *
+ * key: stringified value or key of instance methods
+ * value: count of matching assignments
+ */
+export function getColorAssignmentMatcher(assignments: ColorMapping.Assignment[]) {
+  return new ColorAssignmentMatcher(assignments);
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/color_handling.test.ts

@@ -66,9 +66,11 @@ describe('Color mapping - color generation', () => {
             color: {
               type: 'loop',
             },
-            rule: {
-              type: 'other',
-            },
+            rules: [
+              {
+                type: 'other',
+              },
+            ],
             touched: false,
           },
         ],
@@ -102,9 +104,11 @@ describe('Color mapping - color generation', () => {
               paletteId: KbnPalette.Neutral,
               colorIndex: DEFAULT_NEUTRAL_PALETTE_INDEX,
             },
-            rule: {
-              type: 'other',
-            },
+            rules: [
+              {
+                type: 'other',
+              },
+            ],
             touched: false,
           },
         ],
@@ -143,7 +147,7 @@ describe('Color mapping - color generation', () => {
         assignments: [
           {
             color: { type: 'colorCode', colorCode: 'red' },
-            rule: { type: 'matchExactly', values: ['configuredAssignment'] },
+            rules: [{ type: 'raw', value: 'configuredAssignment' }],
             touched: false,
           },
         ],
@@ -168,17 +172,17 @@ describe('Color mapping - color generation', () => {
         assignments: [
           {
             color: { type: 'colorCode', colorCode: 'red' },
-            rule: { type: 'auto' },
+            rules: [], // auto
             touched: false,
           },
           {
             color: { type: 'colorCode', colorCode: 'blue' },
-            rule: { type: 'matchExactly', values: ['blueCat'] },
+            rules: [{ type: 'raw', value: 'blueCat' }],
             touched: false,
           },
           {
             color: { type: 'colorCode', colorCode: 'green' },
-            rule: { type: 'auto' },
+            rules: [], // auto
             touched: false,
           },
         ],
@@ -267,8 +271,8 @@ describe('Color mapping - color generation', () => {
       {
         ...DEFAULT_COLOR_MAPPING_CONFIG,
         assignments: [
-          { color: { type: 'gradient' }, rule: { type: 'auto' }, touched: false },
-          { color: { type: 'gradient' }, rule: { type: 'auto' }, touched: false },
+          { color: { type: 'gradient' }, rules: [], touched: false },
+          { color: { type: 'gradient' }, rules: [], touched: false },
         ],
 
         colorMode: {
@@ -290,9 +294,11 @@ describe('Color mapping - color generation', () => {
               colorIndex: DEFAULT_NEUTRAL_PALETTE_INDEX,
               paletteId: KbnPalette.Neutral,
             },
-            rule: {
-              type: 'other',
-            },
+            rules: [
+              {
+                type: 'other',
+              },
+            ],
             touched: false,
           },
         ],

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/color_handling.ts

@@ -8,22 +8,25 @@
  */
 
 import chroma from 'chroma-js';
-import { findLast } from 'lodash';
 import { KbnPalette, KbnPalettes } from '@kbn/palettes';
+import { RawValue, SerializedValue, deserializeField } from '@kbn/data-plugin/common';
 import { ColorMapping } from '../config';
 import { changeAlpha, combineColors, getValidColor } from './color_math';
 import { ColorMappingInputData } from '../categorical_color_mapping';
-import { ruleMatch } from './rule_matching';
 import { GradientColorMode } from '../config/types';
 import {
   DEFAULT_NEUTRAL_PALETTE_INDEX,
   DEFAULT_OTHER_ASSIGNMENT_INDEX,
 } from '../config/default_color_mapping';
+import { getColorAssignmentMatcher } from './color_assignment_matcher';
+import { getValueKey } from './utils';
+
+const FALLBACK_ASSIGNMENT_COLOR = 'red';
 
 export function getAssignmentColor(
   colorMode: ColorMapping.Config['colorMode'],
   color:
-    | ColorMapping.Config['assignments'][number]['color']
+    | ColorMapping.Assignment['color']
     | (ColorMapping.LoopColor & { paletteId: string; colorIndex: number }),
   palettes: KbnPalettes,
   isDarkMode: boolean,
@@ -37,11 +40,17 @@ export function getAssignmentColor(
       return getColor(color, palettes);
     case 'gradient': {
       if (colorMode.type === 'categorical') {
-        return 'red';
+        return FALLBACK_ASSIGNMENT_COLOR;
       }
       const colorScale = getGradientColorScale(colorMode, palettes, isDarkMode);
-      return total === 0 ? 'red' : total === 1 ? colorScale(0) : colorScale(index / (total - 1));
+      return total === 0
+        ? FALLBACK_ASSIGNMENT_COLOR
+        : total === 1
+        ? colorScale(0)
+        : colorScale(index / (total - 1));
     }
+    default:
+      return FALLBACK_ASSIGNMENT_COLOR;
   }
 }
 
@@ -57,86 +66,94 @@ export function getColor(
     : getValidColor(palettes.get(color.paletteId).getColor(color.colorIndex)).hex();
 }
 
+/**
+ * Returns a color given a raw value
+ */
+export type ColorHandlingFn = (rawValue: RawValue) => string;
+
 export function getColorFactory(
   { assignments, specialAssignments, colorMode, paletteId }: ColorMapping.Config,
   palettes: KbnPalettes,
   isDarkMode: boolean,
   data: ColorMappingInputData
-): (category: string | string[]) => string {
-  // find auto-assigned colors
-  const autoByOrderAssignments =
-    data.type === 'categories'
-      ? assignments.filter((a) => {
-          return (
-            a.rule.type === 'auto' || (a.rule.type === 'matchExactly' && a.rule.values.length === 0)
-          );
-        })
-      : [];
-
-  // find all categories that don't match with an assignment
-  const notAssignedCategories =
-    data.type === 'categories'
-      ? data.categories.filter((category) => {
-          return !assignments.some(({ rule }) => ruleMatch(rule, category));
-        })
-      : [];
-
-  const lastCategorical = findLast(assignments, (d) => {
-    return d.color.type === 'categorical';
-  });
+): ColorHandlingFn {
+  const lastCategorical = assignments.findLast(({ color }) => color.type === 'categorical');
   const nextCategoricalIndex =
     lastCategorical?.color.type === 'categorical' ? lastCategorical.color.colorIndex + 1 : 0;
 
-  return (category: string | string[]) => {
-    if (typeof category === 'string' || Array.isArray(category)) {
-      const nonAssignedCategoryIndex = notAssignedCategories.indexOf(category);
+  const autoAssignments = assignments
+    .filter(({ rules }) => rules.length === 0)
+    .map((assignment, i) => ({
+      assignment,
+      assignmentIndex: i,
+    }));
 
-      // this category is not assigned to a specific color
-      if (nonAssignedCategoryIndex > -1) {
-        // if the category order is within current number of auto-assigned items pick the defined color
-        if (nonAssignedCategoryIndex < autoByOrderAssignments.length) {
-          const autoAssignmentIndex = assignments.findIndex(
-            (d) => d === autoByOrderAssignments[nonAssignedCategoryIndex]
-          );
-          return getAssignmentColor(
-            colorMode,
-            autoByOrderAssignments[nonAssignedCategoryIndex].color,
-            palettes,
-            isDarkMode,
-            autoAssignmentIndex,
-            assignments.length
-          );
-        }
-        const totalColorsIfGradient = assignments.length || notAssignedCategories.length;
-        const indexIfGradient =
-          (nonAssignedCategoryIndex - autoByOrderAssignments.length) % totalColorsIfGradient;
+  const assignmentMatcher = getColorAssignmentMatcher(assignments);
+  // find all categories that don't match with an assignment
+  const unassignedAutoAssignmentsMap = new Map(
+    data.type === 'categories'
+      ? data.categories
+          .map((category: SerializedValue) => deserializeField(category))
+          .filter((category: RawValue) => {
+            return !assignmentMatcher.hasMatch(category);
+          })
+          .map((category: RawValue, i) => {
+            const key = getValueKey(category);
+            const autoAssignment = autoAssignments[i];
+            return [key, { ...autoAssignment, categoryIndex: i }];
+          })
+      : []
+  );
 
-        // if no auto-assign color rule/color is available then use the color looping palette
+  return (rawValue: RawValue) => {
+    const key = getValueKey(rawValue);
+
+    if (unassignedAutoAssignmentsMap.has(key)) {
+      const {
+        assignment,
+        assignmentIndex = -1,
+        categoryIndex = -1,
+      } = unassignedAutoAssignmentsMap.get(key) ?? {};
+
+      if (assignment) {
+        // the category is within the number of available auto-assignments
         return getAssignmentColor(
           colorMode,
-          // TODO: the specialAssignment[0] position is arbitrary, we should fix it better
-          specialAssignments[DEFAULT_OTHER_ASSIGNMENT_INDEX].color.type === 'loop'
-            ? colorMode.type === 'gradient'
-              ? { type: 'gradient' }
-              : {
-                  type: 'loop',
-                  // those are applied here and depends on the current non-assigned category - auto-assignment list
-                  colorIndex:
-                    nonAssignedCategoryIndex - autoByOrderAssignments.length + nextCategoricalIndex,
-                  paletteId,
-                }
-            : specialAssignments[DEFAULT_OTHER_ASSIGNMENT_INDEX].color,
+          assignment.color,
           palettes,
           isDarkMode,
-          indexIfGradient,
-          totalColorsIfGradient
+          assignmentIndex,
+          assignments.length
         );
       }
+
+      // the category is not assigned to a specific color
+      const totalColorsIfGradient = assignments.length || unassignedAutoAssignmentsMap.size;
+      const indexIfGradient = (categoryIndex - autoAssignments.length) % totalColorsIfGradient;
+
+      // if no auto-assign color rule/color is available then use the color looping palette
+      return getAssignmentColor(
+        colorMode,
+        // TODO: the specialAssignment[0] position is arbitrary, we should fix it better
+        specialAssignments[DEFAULT_OTHER_ASSIGNMENT_INDEX].color.type === 'loop'
+          ? colorMode.type === 'gradient'
+            ? { type: 'gradient' }
+            : {
+                type: 'loop',
+                // those are applied here and depends on the current non-assigned category - auto-assignment list
+                colorIndex: categoryIndex - autoAssignments.length + nextCategoricalIndex,
+                paletteId,
+              }
+          : specialAssignments[DEFAULT_OTHER_ASSIGNMENT_INDEX].color,
+        palettes,
+        isDarkMode,
+        indexIfGradient,
+        totalColorsIfGradient
+      );
     }
+
     // find the assignment where the category matches the rule
-    const matchingAssignmentIndex = assignments.findIndex(({ rule }) => {
-      return ruleMatch(rule, category);
-    });
+    const matchingAssignmentIndex = assignmentMatcher.getIndex(rawValue);
 
     if (matchingAssignmentIndex > -1) {
       const assignment = assignments[matchingAssignmentIndex];

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/rule_matching.ts

@@ -1,64 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the "Elastic License
- * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-
-import { i18n } from '@kbn/i18n';
-import { ColorMapping } from '../config';
-
-export function ruleMatch(
-  rule: ColorMapping.Config['assignments'][number]['rule'],
-  value: string | number | string[]
-) {
-  switch (rule.type) {
-    case 'matchExactly':
-      if (Array.isArray(value)) {
-        return rule.values.some(
-          (v) =>
-            Array.isArray(v) && v.length === value.length && v.every((part, i) => part === value[i])
-        );
-      }
-      return rule.values.includes(`${value}`);
-    case 'matchExactlyCI':
-      return rule.values.some((d) => d.toLowerCase() === `${value}`.toLowerCase());
-    case 'range':
-      // TODO: color by value not yet possible in all charts in elastic-charts
-      return typeof value === 'number' ? rangeMatch(rule, value) : false;
-    default:
-      return false;
-  }
-}
-
-export function rangeMatch(rule: ColorMapping.RuleRange, value: number) {
-  return (
-    (rule.min === rule.max && rule.min === value) ||
-    ((rule.minInclusive ? value >= rule.min : value > rule.min) &&
-      (rule.maxInclusive ? value <= rule.max : value < rule.max))
-  );
-}
-
-// TODO: move in some data/table related package
-export const SPECIAL_TOKENS_STRING_CONVERSION = new Map([
-  [
-    '__other__',
-    i18n.translate('coloring.colorMapping.terms.otherBucketLabel', {
-      defaultMessage: 'Other',
-    }),
-  ],
-  [
-    '',
-    i18n.translate('coloring.colorMapping.terms.emptyLabel', {
-      defaultMessage: '(empty)',
-    }),
-  ],
-]);
-
-/**
- * Returns special string for sake of color mapping/syncing
- */
-export const getSpecialString = (value: string) =>
-  SPECIAL_TOKENS_STRING_CONVERSION.get(value) ?? value;

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/utils.test.ts

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+/* eslint-disable max-classes-per-file */
+
+import { getValueKey } from './utils';
+
+describe('Utils', () => {
+  describe('getValueKey', () => {
+    class WithoutToString {
+      value = 'without toString';
+    }
+    class WithToString {
+      value = 'with toString';
+      toString() {
+        return this.value;
+      }
+    }
+
+    it.each<[desc: string, rawValue: unknown, expectedKey: string]>([
+      ['object', { test: 'test' }, '{"test":"test"}'],
+      ['array', [1, 2, 'three'], '1,2,three'],
+      ['number', 123, '123'],
+      ['string', 'testing', 'testing'],
+      ['boolean (false)', false, 'false'],
+      ['boolean (true)', true, 'true'],
+      ['null', null, 'null'],
+      ['undefined', undefined, 'undefined'],
+      ['class (with toString)', new WithToString(), 'with toString'],
+      ['class (without toString)', new WithoutToString(), '{"value":"without toString"}'],
+    ])('should return correct key for %s', (_, rawValue, expectedKey) => {
+      expect(getValueKey(rawValue)).toBe(expectedKey);
+    });
+  });
+});

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/color/utils.ts

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { RawValue } from '@kbn/data-plugin/common';
+
+/**
+ * Returns string key given an unknown raw color assignment value
+ */
+export function getValueKey(rawValue: RawValue): string {
+  const key = String(rawValue);
+  return key !== '[object Object]' ? key : JSON.stringify(rawValue);
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/assignment/assignment.tsx

@@ -14,10 +14,11 @@ import { i18n } from '@kbn/i18n';
 import { css } from '@emotion/react';
 import { euiThemeVars } from '@kbn/ui-theme';
 import { IKbnPalette, KbnPalettes } from '@kbn/palettes';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
 import {
   removeAssignment,
   updateAssignmentColor,
-  updateAssignmentRule,
+  updateAssignmentRules,
 } from '../../state/color_mapping';
 import { ColorMapping } from '../../config';
 import { Range } from './range';
@@ -25,31 +26,36 @@ import { Match } from './match';
 
 import { ColorMappingInputData } from '../../categorical_color_mapping';
 import { ColorSwatch } from '../color_picker/color_swatch';
+import { ColorAssignmentMatcher } from '../../color/color_assignment_matcher';
 
 export function Assignment({
   data,
   assignment,
+  assignments,
   disableDelete,
   index,
-  total,
   palette,
   palettes,
   colorMode,
   isDarkMode,
   specialTokens,
-  assignmentValuesCounter,
+  formatter,
+  allowCustomMatch,
+  assignmentMatcher,
 }: {
   data: ColorMappingInputData;
   index: number;
-  total: number;
   colorMode: ColorMapping.Config['colorMode'];
-  assignment: ColorMapping.Config['assignments'][number];
+  assignment: ColorMapping.Assignment;
+  assignments: ColorMapping.Assignment[];
   disableDelete: boolean;
   palette: IKbnPalette;
   palettes: KbnPalettes;
   isDarkMode: boolean;
   specialTokens: Map<string, string>;
-  assignmentValuesCounter: Map<string | string[], number>;
+  formatter?: IFieldFormat;
+  allowCustomMatch?: boolean;
+  assignmentMatcher: ColorAssignmentMatcher;
 }) {
   const dispatch = useDispatch();
 
@@ -65,34 +71,21 @@ export function Assignment({
           index={index}
           palette={palette}
           palettes={palettes}
-          total={total}
+          total={assignments.length}
           onColorChange={(color) => {
-            dispatch(updateAssignmentColor({ assignmentIndex: index, color }));
+            dispatch(
+              updateAssignmentColor({
+                assignmentIndex: index,
+                color,
+              })
+            );
           }}
         />
       </EuiFlexItem>
 
-      {assignment.rule.type === 'auto' ||
-      assignment.rule.type === 'matchExactly' ||
-      assignment.rule.type === 'matchExactlyCI' ? (
-        <Match
-          index={index}
-          rule={assignment.rule}
-          options={data.type === 'categories' ? data.categories : []}
-          specialTokens={specialTokens}
-          updateValue={(values: Array<string | string[]>) => {
-            dispatch(
-              updateAssignmentRule({
-                assignmentIndex: index,
-                rule: values.length === 0 ? { type: 'auto' } : { type: 'matchExactly', values },
-              })
-            );
-          }}
-          assignmentValuesCounter={assignmentValuesCounter}
-        />
-      ) : assignment.rule.type === 'range' ? (
+      {assignment.rules[0]?.type === 'range' ? (
         <Range
-          rule={assignment.rule}
+          rule={assignment.rules[0]}
           updateValue={(min, max, minInclusive, maxInclusive) => {
             const rule: ColorMapping.RuleRange = {
               type: 'range',
@@ -101,10 +94,33 @@ export function Assignment({
               minInclusive,
               maxInclusive,
             };
-            dispatch(updateAssignmentRule({ assignmentIndex: index, rule }));
+            dispatch(
+              updateAssignmentRules({
+                assignmentIndex: index,
+                rules: [rule],
+              })
+            );
           }}
         />
-      ) : null}
+      ) : (
+        <Match
+          index={index}
+          rules={assignment.rules}
+          categories={data.type === 'categories' ? data.categories : []}
+          specialTokens={specialTokens}
+          formatter={formatter}
+          allowCustomMatch={allowCustomMatch}
+          assignmentMatcher={assignmentMatcher}
+          updateRules={(rules) => {
+            dispatch(
+              updateAssignmentRules({
+                assignmentIndex: index,
+                rules,
+              })
+            );
+          }}
+        />
+      )}
 
       <EuiFlexItem grow={0}>
         <EuiButtonIcon

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/assignment/duplicate_warning.tsx

@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import React from 'react';
+import { EuiIcon, EuiToolTip } from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { euiThemeVars } from '@kbn/ui-theme';
+
+const duplicateWarning = i18n.translate(
+  'coloring.colorMapping.assignments.duplicateCategoryWarning',
+  {
+    defaultMessage:
+      'This category has already been assigned a different color. Only the first matching assignment will be used.',
+  }
+);
+
+export function DuplicateWarning() {
+  return (
+    <EuiToolTip position="bottom" content={duplicateWarning}>
+      <EuiIcon size="s" type="warning" color={euiThemeVars.euiColorWarningText} />
+    </EuiToolTip>
+  );
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/assignment/match.tsx

@@ -8,68 +8,88 @@
  */
 
 import React from 'react';
-import { EuiComboBox, EuiFlexItem, EuiIcon, EuiToolTip } from '@elastic/eui';
+import { EuiComboBox, EuiComboBoxOptionOption, EuiFlexItem } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
-import { MULTI_FIELD_KEY_SEPARATOR } from '@kbn/data-plugin/common';
-import { euiThemeVars } from '@kbn/ui-theme';
+import { RawValue, SerializedValue, deserializeField } from '@kbn/data-plugin/common';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
 import { ColorMapping } from '../../config';
+import { ColorRule, RuleMatch, RuleMatchRaw } from '../../config/types';
+import { ColorAssignmentMatcher } from '../../color/color_assignment_matcher';
+import { DuplicateWarning } from './duplicate_warning';
+import { getValueKey } from '../../color/utils';
+
+export const isNotNull = <T,>(value: T | null): value is NonNullable<T> => value !== null;
 
 export const Match: React.FC<{
   index: number;
-  rule:
-    | ColorMapping.RuleAuto
-    | ColorMapping.RuleMatchExactly
-    | ColorMapping.RuleMatchExactlyCI
-    | ColorMapping.RuleRegExp;
-  updateValue: (values: Array<string | string[]>) => void;
-  options: Array<string | string[]>;
+  rules: ColorMapping.ColorRule[];
+  updateRules: (rule: ColorMapping.ColorRule[]) => void;
+  categories: SerializedValue[];
   specialTokens: Map<unknown, string>;
-  assignmentValuesCounter: Map<string | string[], number>;
-}> = ({ index, rule, updateValue, options, specialTokens, assignmentValuesCounter }) => {
-  const duplicateWarning = i18n.translate(
-    'coloring.colorMapping.assignments.duplicateCategoryWarning',
-    {
-      defaultMessage:
-        'This category has already been assigned a different color. Only the first matching assignment will be used.',
-    }
-  );
-  const selectedOptions =
-    rule.type === 'auto'
-      ? []
-      : typeof rule.values === 'string'
-      ? [
-          {
-            label: rule.values,
-            value: rule.values,
-            append:
-              (assignmentValuesCounter.get(rule.values) ?? 0) > 1 ? (
-                <EuiToolTip position="bottom" content={duplicateWarning}>
-                  <EuiIcon size="s" type="warning" color={euiThemeVars.euiColorWarningText} />
-                </EuiToolTip>
-              ) : undefined,
-          },
-        ]
-      : rule.values.map((value) => {
-          const ruleValues = Array.isArray(value) ? value : [value];
-          return {
-            label: ruleValues.map((v) => specialTokens.get(v) ?? v).join(MULTI_FIELD_KEY_SEPARATOR),
-            value,
-            append:
-              (assignmentValuesCounter.get(value) ?? 0) > 1 ? (
-                <EuiToolTip position="bottom" content={duplicateWarning}>
-                  <EuiIcon size="s" type="warning" color={euiThemeVars.euiColorWarningText} />
-                </EuiToolTip>
-              ) : undefined,
-          };
-        });
+  formatter?: IFieldFormat;
+  allowCustomMatch?: boolean;
+  assignmentMatcher: ColorAssignmentMatcher;
+}> = ({
+  index,
+  rules,
+  updateRules,
+  categories,
+  specialTokens,
+  formatter,
+  allowCustomMatch = false,
+  assignmentMatcher,
+}) => {
+  const getOptionForRawValue = getOptionForRawValueFn(formatter);
+  const availableOptions: Array<EuiComboBoxOptionOption<string>> = [];
 
-  const convertedOptions = options.map((value) => {
-    const ruleValues = Array.isArray(value) ? value : [value];
-    return {
-      label: ruleValues.map((v) => specialTokens.get(v) ?? v).join(MULTI_FIELD_KEY_SEPARATOR),
-      value,
-    };
-  });
+  // Map option key to their raw value
+  const rawCategoryValueMap = categories.reduce<Map<string, RawValue>>(
+    (acc, value: SerializedValue) => {
+      const option = getOptionForRawValue(value);
+      availableOptions.push(option);
+      acc.set(option.key, value);
+      return acc;
+    },
+    new Map()
+  );
+  const selectedOptions = rules
+    .map<EuiComboBoxOptionOption<string> | null>((rule) => {
+      switch (rule.type) {
+        case 'raw': {
+          const rawValue = deserializeField(rule.value);
+          const hasDuplicate = assignmentMatcher.hasDuplicate(rawValue);
+          const option = getOptionForRawValue(rule.value);
+          rawCategoryValueMap.set(option.key, rule.value);
+          return {
+            ...option,
+            append: hasDuplicate && <DuplicateWarning />,
+          };
+        }
+        case 'match': {
+          const key = rule.matchCase ? rule.pattern : rule.pattern.toLowerCase();
+          const hasDuplicate = assignmentMatcher.hasDuplicate(key); // non-exhaustive for partial word match
+
+          return {
+            label: specialTokens.get(rule.pattern) ?? rule.pattern,
+            key: rule.pattern,
+            append: hasDuplicate && <DuplicateWarning />,
+          };
+        }
+        case 'regex': {
+          // Note: Only basic placeholder logic, not used or fully tested
+          const hasDuplicate = false; // need to use exhaustive search
+
+          return {
+            label: rule.pattern,
+            key: rule.pattern,
+            append: hasDuplicate && <DuplicateWarning />,
+          };
+        }
+        default:
+          return null;
+      }
+    })
+    .filter(isNotNull);
 
   return (
     <EuiFlexItem style={{ minWidth: 1, width: 1 }}>
@@ -87,26 +107,67 @@ export const Match: React.FC<{
             defaultMessage: 'Auto assigning term',
           }
         )}
-        options={convertedOptions}
+        options={availableOptions}
         selectedOptions={selectedOptions}
-        onChange={(changedOptions) => {
-          updateValue(
-            changedOptions.reduce<Array<string | string[]>>((acc, option) => {
-              if (option.value !== undefined) {
-                acc.push(option.value);
+        onChange={(newOptions) => {
+          updateRules(
+            newOptions.reduce<ColorRule[]>((acc, { key = null }) => {
+              if (key !== null) {
+                if (rawCategoryValueMap.has(key)) {
+                  acc.push({
+                    type: 'raw',
+                    value: rawCategoryValueMap.get(key),
+                  } satisfies RuleMatchRaw);
+                } else {
+                  acc.push({
+                    type: 'match',
+                    pattern: key,
+                    matchCase: true,
+                    matchEntireWord: true,
+                  } satisfies RuleMatch);
+                }
               }
               return acc;
             }, [])
           );
         }}
-        onCreateOption={(label) => {
-          if (selectedOptions.findIndex((option) => option.label === label) === -1) {
-            updateValue([...selectedOptions, { label, value: label }].map((d) => d.value));
-          }
+        optionMatcher={({ option, normalizedSearchValue }) => {
+          return (
+            String(option.value ?? '').includes(normalizedSearchValue) ||
+            option.label.includes(normalizedSearchValue)
+          );
         }}
+        onCreateOption={
+          allowCustomMatch
+            ? (label) => {
+                return updateRules([
+                  ...rules,
+                  {
+                    type: 'match',
+                    pattern: label,
+                    matchCase: true,
+                    matchEntireWord: true,
+                  } satisfies RuleMatch,
+                ]);
+              }
+            : undefined
+        }
         isCaseSensitive
         compressed
       />
     </EuiFlexItem>
   );
 };
+
+function getOptionForRawValueFn(fieldFormat?: IFieldFormat) {
+  const formatter = fieldFormat?.convert.bind(fieldFormat) ?? String;
+  return (serializedValue: unknown) => {
+    const rawValue = deserializeField(serializedValue);
+    const key = getValueKey(rawValue);
+    return {
+      key,
+      value: typeof rawValue === 'number' ? key : undefined,
+      label: formatter(rawValue),
+    } satisfies EuiComboBoxOptionOption<string>;
+  };
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/assignment/utils.ts

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { DatatableColumnMeta } from '@kbn/expressions-plugin/common';
+
+/**
+ * Determines is a field can have create a custom pattern to match during color mapping.
+ */
+export function canCreateCustomMatch(meta?: DatatableColumnMeta): boolean {
+  if (!meta) return false;
+  return (
+    (meta.type === 'number' || meta.type === 'string') &&
+    meta.params?.id !== 'range' &&
+    meta.params?.id !== 'multi_terms'
+  );
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/color_picker/color_picker.tsx

@@ -46,12 +46,20 @@ export function ColorPicker({
         }}
       >
         <EuiTabs size="m" expand>
-          <EuiTab onClick={() => setTab('palette')} isSelected={tab === 'palette'}>
+          <EuiTab
+            data-test-subj="lns-colorMapping-colorPicker-tab-colors"
+            onClick={() => setTab('palette')}
+            isSelected={tab === 'palette'}
+          >
             {i18n.translate('coloring.colorMapping.colorPicker.paletteTabLabel', {
               defaultMessage: 'Colors',
             })}
           </EuiTab>
-          <EuiTab onClick={() => setTab('custom')} isSelected={tab === 'custom'}>
+          <EuiTab
+            data-test-subj="lns-colorMapping-colorPicker-tab-custom"
+            onClick={() => setTab('custom')}
+            isSelected={tab === 'custom'}
+          >
             {i18n.translate('coloring.colorMapping.colorPicker.customTabLabel', {
               defaultMessage: 'Custom',
             })}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/color_picker/color_swatch.tsx

@@ -30,7 +30,7 @@ import { getValidColor } from '../../color/color_math';
 
 interface ColorPickerSwatchProps {
   colorMode: ColorMapping.Config['colorMode'];
-  assignmentColor: ColorMapping.Config['assignments'][number]['color'];
+  assignmentColor: ColorMapping.Assignment['color'];
   index: number;
   total: number;
   palette: IKbnPalette;
@@ -163,8 +163,11 @@ export const ColorSwatch = ({
   ) : (
     <EuiColorPickerSwatch
       color={colorHex}
-      aria-label={i18n.translate('coloring.colorMapping.colorPicker.newColorAriaLabel', {
-        defaultMessage: 'Select a new color',
+      aria-label={i18n.translate('coloring.colorMapping.colorSwatch.gradientAriaLabel', {
+        defaultMessage: 'Computed gradient color',
+      })}
+      title={i18n.translate('coloring.colorMapping.colorSwatch.gradientAriaLabel', {
+        defaultMessage: 'Computed gradient color',
       })}
       disabled
       style={{

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/color_picker/rgb_picker.tsx

@@ -80,7 +80,7 @@ export function RGBPicker({
     ? euiThemeVars.euiColorWarningText
     : '';
   return (
-    <EuiFlexGroup direction="column" gutterSize="s" style={{ padding: 8 }}>
+    <EuiFlexGroup direction="column" gutterSize="s" css={{ padding: 8 }}>
       <EuiFlexItem>
         <EuiColorPicker
           onChange={(c) => {
@@ -139,6 +139,7 @@ export function RGBPicker({
                       });
                     }
                   }}
+                  data-test-subj="lns-colorMapping-colorPicker-custom-input"
                   aria-label={i18n.translate(
                     'coloring.colorMapping.colorPicker.hexColorinputAriaLabel',
                     {

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/container/assignments.tsx

@@ -30,6 +30,8 @@ import { i18n } from '@kbn/i18n';
 import { useDispatch, useSelector } from 'react-redux';
 import { findLast } from 'lodash';
 import { KbnPalettes } from '@kbn/palettes';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
+import { deserializeField } from '@kbn/data-plugin/common';
 import { Assignment } from '../assignment/assignment';
 import {
   addNewAssignment,
@@ -39,19 +41,25 @@ import {
 import { selectColorMode, selectComputedAssignments, selectPalette } from '../../state/selectors';
 import { ColorMappingInputData } from '../../categorical_color_mapping';
 import { ColorMapping } from '../../config';
-import { ruleMatch } from '../../color/rule_matching';
+import { getColorAssignmentMatcher } from '../../color/color_assignment_matcher';
 
-export function AssignmentsConfig({
+export function Assignments({
   data,
   palettes,
   isDarkMode,
   specialTokens,
+  formatter,
+  allowCustomMatch,
 }: {
   palettes: KbnPalettes;
   data: ColorMappingInputData;
   isDarkMode: boolean;
-  /** map between original and formatted tokens used to handle special cases, like the Other bucket and the empty bucket */
+  /**
+   * map between original and formatted tokens used to handle special cases, like the Other bucket and the empty bucket
+   */
   specialTokens: Map<string, string>;
+  formatter?: IFieldFormat;
+  allowCustomMatch?: boolean;
 }) {
   const [showOtherActions, setShowOtherActions] = useState<boolean>(false);
 
@@ -59,41 +67,25 @@ export function AssignmentsConfig({
   const palette = useSelector(selectPalette(palettes));
   const colorMode = useSelector(selectColorMode);
   const assignments = useSelector(selectComputedAssignments);
-
+  const assignmentMatcher = useMemo(() => getColorAssignmentMatcher(assignments), [assignments]);
   const unmatchingCategories = useMemo(() => {
     return data.type === 'categories'
       ? data.categories.filter((category) => {
-          return !assignments.some(({ rule }) => ruleMatch(rule, category));
+          const rawValue = deserializeField(category);
+          return !assignmentMatcher.hasMatch(rawValue);
         })
       : [];
-  }, [data, assignments]);
-
-  const assignmentValuesCounter = assignments.reduce<Map<string | string[], number>>(
-    (acc, assignment) => {
-      const values = assignment.rule.type === 'matchExactly' ? assignment.rule.values : [];
-      values.forEach((value) => {
-        acc.set(value, (acc.get(value) ?? 0) + 1);
-      });
-      return acc;
-    },
-    new Map()
-  );
+  }, [data, assignmentMatcher]);
 
   const onClickAddNewAssignment = useCallback(() => {
-    const lastCategorical = findLast(assignments, (d) => {
-      return d.color.type === 'categorical';
+    const lastCategorical = assignments.findLast((a) => {
+      return a.color.type === 'categorical';
     });
     const nextCategoricalIndex =
       lastCategorical?.color.type === 'categorical' ? lastCategorical.color.colorIndex + 1 : 0;
     dispatch(
       addNewAssignment({
-        rule:
-          data.type === 'categories'
-            ? {
-                type: 'matchExactly',
-                values: [],
-              }
-            : { type: 'range', min: 0, max: 0, minInclusive: true, maxInclusive: true },
+        rules: [],
         color:
           colorMode.type === 'categorical'
             ? {
@@ -105,7 +97,7 @@ export function AssignmentsConfig({
         touched: false,
       })
     );
-  }, [assignments, colorMode.type, data.type, dispatch, palette]);
+  }, [assignments, colorMode.type, dispatch, palette]);
 
   const onClickAddAllCurrentCategories = useCallback(() => {
     if (data.type === 'categories') {
@@ -115,25 +107,25 @@ export function AssignmentsConfig({
       const nextCategoricalIndex =
         lastCategorical?.color.type === 'categorical' ? lastCategorical.color.colorIndex + 1 : 0;
 
-      const newAssignments: ColorMapping.Config['assignments'] = unmatchingCategories.map(
-        (c, i) => {
-          return {
-            rule: {
-              type: 'matchExactly',
-              values: [c],
+      const newAssignments = unmatchingCategories.map((category, i) => {
+        return {
+          rules: [
+            {
+              type: 'raw',
+              value: category,
             },
-            color:
-              colorMode.type === 'categorical'
-                ? {
-                    type: 'categorical',
-                    paletteId: palette.id,
-                    colorIndex: (nextCategoricalIndex + i) % palette.colors().length,
-                  }
-                : { type: 'gradient' },
-            touched: false,
-          };
-        }
-      );
+          ],
+          color:
+            colorMode.type === 'categorical'
+              ? {
+                  type: 'categorical',
+                  paletteId: palette.id,
+                  colorIndex: (nextCategoricalIndex + i) % palette.colors().length,
+                }
+              : { type: 'gradient' },
+          touched: false,
+        } satisfies ColorMapping.Assignment;
+      });
       dispatch(addNewAssignments(newAssignments));
     }
   }, [data.type, assignments, unmatchingCategories, dispatch, colorMode.type, palette]);
@@ -164,7 +156,7 @@ export function AssignmentsConfig({
                 key={i}
                 data={data}
                 index={i}
-                total={assignments.length}
+                assignments={assignments}
                 colorMode={colorMode}
                 palette={palette}
                 palettes={palettes}
@@ -172,7 +164,9 @@ export function AssignmentsConfig({
                 assignment={assignment}
                 disableDelete={false}
                 specialTokens={specialTokens}
-                assignmentValuesCounter={assignmentValuesCounter}
+                formatter={formatter}
+                allowCustomMatch={allowCustomMatch}
+                assignmentMatcher={assignmentMatcher}
               />
             );
           })}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/components/container/container.tsx

@@ -13,6 +13,7 @@ import { EuiFlexGroup, EuiFlexItem, EuiFormRow, EuiButtonIcon, EuiToolTip } from
 import { i18n } from '@kbn/i18n';
 import { css } from '@emotion/react';
 import { KbnPalettes } from '@kbn/palettes';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
 import { PaletteSelector } from '../palette_selector/palette_selector';
 
 import { changeGradientSortOrder } from '../../state/color_mapping';
@@ -21,19 +22,23 @@ import { ColorMappingInputData } from '../../categorical_color_mapping';
 import { Gradient } from '../palette_selector/gradient';
 import { ScaleMode } from '../palette_selector/scale';
 import { UnassignedTermsConfig } from './unassigned_terms_config';
-import { AssignmentsConfig } from './assigments';
+import { Assignments } from './assignments';
 
 export function Container({
   data,
   palettes,
   isDarkMode,
   specialTokens,
+  formatter,
+  allowCustomMatch,
 }: {
   palettes: KbnPalettes;
   data: ColorMappingInputData;
   isDarkMode: boolean;
   /** map between original and formatted tokens used to handle special cases, like the Other bucket and the empty bucket */
   specialTokens: Map<string, string>;
+  formatter?: IFieldFormat;
+  allowCustomMatch?: boolean;
 }) {
   const dispatch = useDispatch();
   const palette = useSelector(selectPalette(palettes));
@@ -111,11 +116,13 @@ export function Container({
             defaultMessage: 'Color assignments',
           })}
         >
-          <AssignmentsConfig
+          <Assignments
             isDarkMode={isDarkMode}
             data={data}
             palettes={palettes}
             specialTokens={specialTokens}
+            formatter={formatter}
+            allowCustomMatch={allowCustomMatch}
           />
         </EuiFormRow>
       </EuiFlexItem>

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/assignments.ts

@@ -18,11 +18,11 @@ export function updateAssignmentsPalette(
   preserveColorChanges: boolean
 ): ColorMapping.Config['assignments'] {
   const palette = palettes.get(paletteId);
-  return assignments.map(({ rule, color, touched }, index) => {
+  return assignments.map(({ rules, color, touched }, index) => {
     if (preserveColorChanges && touched) {
-      return { rule, color, touched };
+      return { rules, color, touched };
     } else {
-      const newColor: ColorMapping.Config['assignments'][number]['color'] =
+      const newColor: ColorMapping.Assignment['color'] =
         colorMode.type === 'categorical'
           ? {
               type: 'categorical',
@@ -31,7 +31,7 @@ export function updateAssignmentsPalette(
             }
           : { type: 'gradient' };
       return {
-        rule,
+        rules,
         color: newColor,
         touched: false,
       };

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/colors.ts

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+/**
+ * Defines explicit color specified as a CSS color datatype (rgb/a,hex,keywords,lab,lch etc)
+ */
+export interface ColorCode {
+  type: 'colorCode';
+  colorCode: string;
+}
+
+/**
+ * Defines categorical color based on the index position of color in palette defined by the paletteId
+ */
+export interface CategoricalColor {
+  type: 'categorical';
+  paletteId: string;
+  colorIndex: number;
+}
+
+/**
+ * Defines color based on looping round-robin assignment
+ */
+export interface LoopColor {
+  type: 'loop';
+}
+
+/**
+ * Specify that the Color in an Assignment needs to be taken from a gradient defined in the `Config.colorMode`
+ */
+export interface GradientColor {
+  type: 'gradient';
+}
+
+export type Color = ColorCode | CategoricalColor | LoopColor | GradientColor;

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/default_color_mapping.ts

@@ -22,9 +22,11 @@ export const DEFAULT_COLOR_MAPPING_CONFIG: ColorMapping.Config = {
   assignments: [],
   specialAssignments: [
     {
-      rule: {
-        type: 'other',
-      },
+      rules: [
+        {
+          type: 'other',
+        },
+      ],
       color: {
         type: 'loop',
       },

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/rules.ts

@@ -0,0 +1,92 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { SerializedValue } from '@kbn/data-plugin/common';
+
+/**
+ * A rule that matches based on raw values
+ */
+export interface RuleMatchRaw {
+  type: 'raw';
+  /**
+   * Serialized form of raw row value
+   */
+  value: SerializedValue;
+}
+
+/**
+ * Rule to match values, optionally case-sensitive and entire word
+ */
+export interface RuleMatch {
+  type: 'match';
+  /**
+   * The string to search for
+   */
+  pattern: string;
+  /**
+   * Whether the pattern should match an entire word
+   *
+   * @default false
+   */
+  matchEntireWord?: boolean;
+  /**
+   * Whether the search should be case-sensitive
+   *
+   * @default false
+   */
+  matchCase?: boolean;
+}
+
+/**
+ * Regex rule
+ */
+export interface RuleRegExp {
+  type: 'regex';
+  /**
+   * RegExp pattern as string including flags (i.e. `/[a-z]+/i`)
+   */
+  pattern: string;
+}
+
+/**
+ * Rule for numerical data range assignments
+ */
+export interface RuleRange {
+  type: 'range';
+  /**
+   * The min value of the range
+   */
+  min: number;
+  /**
+   * The max value of the range
+   */
+  max: number;
+  /**
+   * `true` if the range is left-closed (the `min` value is considered within the range), false otherwise (only values that are
+   * greater than the `min` are considered within the range)
+   */
+  minInclusive: boolean;
+  /**
+   * `true` if the range is right-closed (the `max` value is considered within the range), false otherwise (only values less than
+   * the `max` are considered within the range)
+   */
+  maxInclusive: boolean;
+}
+
+/**
+ * A specific catch-everything-else rule
+ */
+export interface RuleOthers {
+  type: 'other';
+}
+
+/**
+ * All available color rules
+ */
+export type ColorRule = RuleMatchRaw | RuleMatch | RuleRegExp | RuleRange;

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/types.ts

@@ -7,146 +7,77 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-/**
- * A color specified as a CSS color datatype (rgb/a,hex,keywords,lab,lch etc)
- */
-export interface ColorCode {
-  type: 'colorCode';
-  colorCode: string;
-}
-
-/**
- * An index specified categorical color, coming from paletteId
- */
-export interface CategoricalColor {
-  type: 'categorical';
-  paletteId: string;
-  colorIndex: number;
-}
-
-/**
- * Specify that the Color in an Assignment needs to be taken from a gradient defined in the `Config.colorMode`
- */
-export interface GradientColor {
-  type: 'gradient';
-}
-
-/**
- * An index specified categorical color, coming from paletteId
- */
-export interface LoopColor {
-  type: 'loop';
-}
-
-/**
- * A special rule that match automatically, in order, all the categories that are not matching a specified rule
- */
-export interface RuleAuto {
-  /* tag */
-  type: 'auto';
-}
-/**
- * A rule that match exactly, case sensitive, with the provided strings
- */
-export interface RuleMatchExactly {
-  /* tag */
-  type: 'matchExactly';
-  values: Array<string | string[]>;
-}
-
-/**
- * A Match rule to match the values case insensitive
- * @ignore not used yet
- */
-export interface RuleMatchExactlyCI {
-  /* tag */
-  type: 'matchExactlyCI';
-  values: string[];
-}
-
-/**
- * A range rule, not used yet, but can be used for numerical data assignments
- */
-export interface RuleRange {
-  /* tag */
-  type: 'range';
-  /**
-   * The min value of the range
-   */
-  min: number;
-  /**
-   * The max value of the range
-   */
-  max: number;
-  /**
-   * `true` if the range is left-closed (the `min` value is considered within the range), false otherwise (only values that are
-   * greater than the `min` are considered within the range)
-   */
-  minInclusive: boolean;
-  /**
-   * `true` if the range is right-closed (the `max` value is considered within the range), false otherwise (only values less than
-   * the `max` are considered within the range)
-   */
-  maxInclusive: boolean;
-}
-/**
- * Regex rule.
- * @ignore not used yet
- */
-export interface RuleRegExp {
-  /* tag */
-  type: 'regex';
-  /**
-   * TODO: not sure how we can store a regexp
-   */
-  values: string;
-}
-
-/**
- * A specific catch-everything-else rule
- */
-export interface RuleOthers {
-  /* tag */
-  type: 'other';
-}
+import { CategoricalColor, Color, ColorCode, GradientColor, LoopColor } from './colors';
+import { ColorRule, RuleOthers } from './rules';
 
 /**
  * An assignment is the connection link between a rule and a color
  */
-export interface Assignment<R, C> {
+export interface AssignmentBase<R extends ColorRule | RuleOthers, C extends Color> {
   /**
    * Describe the rule used to assign the color.
    */
-  rule: R;
+  rules: R[];
   /**
    * The color definition
    */
   color: C;
-
   /**
    * Specify if the color was changed from the original one
-   * TODO: rename
    */
   touched: boolean;
 }
 
+type ColorStep = (CategoricalColor | ColorCode) & {
+  /**
+   * A flag to know when assignment has been edited since last saved
+   */
+  touched: boolean;
+};
+
 export interface CategoricalColorMode {
   type: 'categorical';
 }
+
 export interface GradientColorMode {
   type: 'gradient';
-  steps: Array<(CategoricalColor | ColorCode) & { touched: boolean }>;
+  steps: ColorStep[];
   sort: 'asc' | 'desc';
 }
 
-export interface Config {
+interface BaseConfig {
   paletteId: string;
-  colorMode: CategoricalColorMode | GradientColorMode;
-  assignments: Array<
-    Assignment<
-      RuleAuto | RuleMatchExactly | RuleMatchExactlyCI | RuleRange | RuleRegExp,
-      CategoricalColor | ColorCode | GradientColor
-    >
-  >;
-  specialAssignments: Array<Assignment<RuleOthers, CategoricalColor | ColorCode | LoopColor>>;
+  specialAssignments: Array<AssignmentBase<RuleOthers, CategoricalColor | ColorCode | LoopColor>>;
 }
+
+/**
+ * Gradient color mapping config
+ */
+export interface GradientConfig extends BaseConfig {
+  colorMode: GradientColorMode;
+  assignments: Array<AssignmentBase<ColorRule, GradientColor>>;
+}
+
+/**
+ * Categorical color mapping config
+ */
+export interface CategoricalConfig extends BaseConfig {
+  colorMode: CategoricalColorMode;
+  assignments: Array<AssignmentBase<ColorRule, CategoricalColor | ColorCode>>;
+}
+
+/**
+ * Polymorphic color mapping config
+ *
+ * Merges `GradientConfig` and `CategoricalConfig` for simplicity of type alignment
+ */
+export type Config = BaseConfig & {
+  colorMode: CategoricalColorMode | GradientColorMode;
+  assignments: Array<AssignmentBase<ColorRule, CategoricalColor | ColorCode | GradientColor>>;
+};
+
+export type Assignment = Config['assignments'][number];
+export type SpecialAssignment = BaseConfig['specialAssignments'][number];
+
+export * from './colors';
+export * from './rules';

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/config/utils.ts

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { Config, CategoricalConfig, GradientConfig } from './types';
+
+export function isCategoricalColorConfig(config: Config): config is CategoricalConfig {
+  return config.colorMode.type === 'categorical';
+}
+
+export function isGradientColorConfig(config: Config): config is GradientConfig {
+  return config.colorMode.type === 'gradient';
+}

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/index.ts

@@ -16,10 +16,13 @@ export {
 export type { ColorMappingInputData } from './categorical_color_mapping';
 export type { ColorMapping } from './config';
 export * from './color/color_handling';
-export { SPECIAL_TOKENS_STRING_CONVERSION, getSpecialString } from './color/rule_matching';
+export { getValueKey } from './color/utils';
+export { SPECIAL_TOKENS_STRING_CONVERSION, getSpecialString } from './special_tokens';
+export { type ColorAssignmentMatcher } from './color/color_assignment_matcher';
 export {
   DEFAULT_COLOR_MAPPING_CONFIG,
   DEFAULT_OTHER_ASSIGNMENT_INDEX,
   getPaletteColors,
   getColorsFromMapping,
 } from './config/default_color_mapping';
+export * from './components/assignment/utils';

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/special_tokens.ts

@@ -0,0 +1,32 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { i18n } from '@kbn/i18n';
+
+// TODO: move in some data/table related package
+export const SPECIAL_TOKENS_STRING_CONVERSION = new Map([
+  [
+    '__other__',
+    i18n.translate('coloring.colorMapping.terms.otherBucketLabel', {
+      defaultMessage: 'Other',
+    }),
+  ],
+  [
+    '',
+    i18n.translate('coloring.colorMapping.terms.emptyLabel', {
+      defaultMessage: '(empty)',
+    }),
+  ],
+]);
+
+/**
+ * Returns special string for sake of color mapping/syncing
+ */
+export const getSpecialString = (value: string) =>
+  SPECIAL_TOKENS_STRING_CONVERSION.get(value) ?? value;

src/platform/packages/shared/kbn-coloring/src/shared_components/color_mapping/state/color_mapping.ts

@@ -59,10 +59,7 @@ export const colorMappingSlice = createSlice({
       state.assignments = [];
     },
 
-    addNewAssignment: (
-      state,
-      action: PayloadAction<ColorMapping.Config['assignments'][number]>
-    ) => {
+    addNewAssignment: (state, action: PayloadAction<ColorMapping.Assignment>) => {
       state.assignments.push({ ...action.payload });
     },
     addNewAssignments: (state, action: PayloadAction<ColorMapping.Config['assignments']>) => {
@@ -72,7 +69,7 @@ export const colorMappingSlice = createSlice({
       state,
       action: PayloadAction<{
         assignmentIndex: number;
-        assignment: ColorMapping.Config['assignments'][number];
+        assignment: ColorMapping.Assignment;
       }>
     ) => {
       state.assignments[action.payload.assignmentIndex] = {
@@ -84,19 +81,37 @@ export const colorMappingSlice = createSlice({
       state,
       action: PayloadAction<{
         assignmentIndex: number;
-        rule: ColorMapping.Config['assignments'][number]['rule'];
+        ruleIndex: number;
+        rule: ColorMapping.ColorRule;
+      }>
+    ) => {
+      const assignment = state.assignments[action.payload.assignmentIndex];
+      state.assignments[action.payload.assignmentIndex] = {
+        ...assignment,
+        rules: [
+          ...assignment.rules.slice(0, action.payload.ruleIndex),
+          action.payload.rule,
+          ...assignment.rules.slice(action.payload.ruleIndex + 1),
+        ],
+      };
+    },
+    updateAssignmentRules: (
+      state,
+      action: PayloadAction<{
+        assignmentIndex: number;
+        rules: ColorMapping.ColorRule[];
       }>
     ) => {
       state.assignments[action.payload.assignmentIndex] = {
         ...state.assignments[action.payload.assignmentIndex],
-        rule: action.payload.rule,
+        rules: action.payload.rules,
       };
     },
     updateAssignmentColor: (
       state,
       action: PayloadAction<{
         assignmentIndex: number;
-        color: ColorMapping.Config['assignments'][number]['color'];
+        color: ColorMapping.Assignment['color'];
       }>
     ) => {
       state.assignments[action.payload.assignmentIndex] = {
@@ -219,6 +234,7 @@ export const colorMappingSlice = createSlice({
     },
   },
 });
+
 // Action creators are generated for each case reducer function
 export const {
   updatePalette,
@@ -230,6 +246,7 @@ export const {
   updateAssignmentColor,
   updateSpecialAssignmentColor,
   updateAssignmentRule,
+  updateAssignmentRules,
   removeAssignment,
   removeAllAssignments,
   changeColorMode,

src/platform/packages/shared/kbn-coloring/tsconfig.json

@@ -26,6 +26,8 @@
     "@kbn/ui-theme",
     "@kbn/visualization-utils",
     "@kbn/palettes",
+    "@kbn/field-formats-plugin",
+    "@kbn/expressions-plugin",
   ],
   "exclude": [
     "target/**/*",

src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/colors/color_mapping_accessors.ts

@@ -9,13 +9,12 @@
 
 import { NodeColorAccessor, PATH_KEY } from '@elastic/charts';
 import { decreaseOpacity } from '@kbn/charts-plugin/public';
-import { MultiFieldKey } from '@kbn/data-plugin/common';
 import { getColorFactory } from '@kbn/coloring';
-import { isMultiFieldKey } from '@kbn/data-plugin/common';
+import { MultiFieldKey } from '@kbn/data-plugin/common';
 import { ChartTypes } from '../../../common/types';
 
 export function getCategoryKeys(category: string | MultiFieldKey): string | string[] {
-  return isMultiFieldKey(category) ? category.keys.map(String) : `${category}`;
+  return MultiFieldKey.isInstance(category) ? category.keys.map(String) : `${category}`;
 }
 
 /**

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

@@ -99,12 +99,12 @@ describe('get color', () => {
     const buckets = createMockBucketColumns();
     const visParams = createMockPieParams();
     const colors = ['color1', 'color2', 'color3', 'color4'];
-    const categories = (chartType?: ChartTypes) =>
+    const getCategories = (chartType?: ChartTypes) =>
       chartType === ChartTypes.MOSAIC && visData.columns.length === 2
         ? getColorCategories(visData.rows, visData.columns[1]?.id)
         : getColorCategories(visData.rows, visData.columns[0]?.id);
-    const colorIndexMap = (chartType?: ChartTypes) =>
-      new Map(categories(chartType).map((d, i) => [d[0], i]));
+    const getColorIndexMap = (chartType?: ChartTypes) =>
+      new Map(getCategories(chartType).map((d, i) => [d, i]));
     const dataMock = dataPluginMock.createStartContract();
     interface RangeProps {
       gte: number;
@@ -146,7 +146,8 @@ describe('get color', () => {
         getAll: () => [mockPalette1],
       };
     };
-    it('should return the correct color based on the parent sortIndex', () => {
+
+    it('should return the correct color based on color map index', () => {
       const d: SimplifiedArrayNode = {
         depth: 1,
         sortIndex: 0,
@@ -177,7 +178,43 @@ describe('get color', () => {
         false,
         dataMock.fieldFormats,
         visData.columns[0],
-        colorIndexMap(ChartTypes.PIE)
+        getColorIndexMap(ChartTypes.PIE)
+      );
+      expect(color).toEqual(colors[2]);
+    });
+
+    it('should return the correct color based on the parent sortIndex when no color map index found', () => {
+      const d: SimplifiedArrayNode = {
+        depth: 1,
+        sortIndex: 0,
+        parent: {
+          children: [
+            ['ES-Air', undefined],
+            ['Kibana Airlines', undefined],
+          ],
+          depth: 0,
+          sortIndex: 0,
+        },
+        children: [],
+      };
+
+      const color = getColor(
+        ChartTypes.PIE,
+        'ES-Air',
+        d,
+        0,
+        false,
+        {},
+        distinctSeries,
+        dataLength,
+        visParams,
+        getPaletteRegistry(),
+        { getColor: () => undefined },
+        false,
+        false,
+        dataMock.fieldFormats,
+        visData.columns[0],
+        new Map()
       );
       expect(color).toEqual(colors[0]);
     });
@@ -212,7 +249,7 @@ describe('get color', () => {
         false,
         dataMock.fieldFormats,
         visData.columns[0],
-        colorIndexMap(ChartTypes.PIE)
+        getColorIndexMap(ChartTypes.PIE)
       );
       expect(color).toEqual('color3');
     });
@@ -246,7 +283,7 @@ describe('get color', () => {
         false,
         dataMock.fieldFormats,
         visData.columns[0],
-        colorIndexMap(ChartTypes.PIE)
+        getColorIndexMap(ChartTypes.PIE)
       );
       expect(color).toEqual('#000028');
     });
@@ -310,7 +347,7 @@ describe('get color', () => {
         false,
         dataMock.fieldFormats,
         column,
-        colorIndexMap(ChartTypes.PIE)
+        getColorIndexMap(ChartTypes.PIE)
       );
       expect(color).toEqual('#3F6833');
     });
@@ -351,7 +388,7 @@ describe('get color', () => {
         false,
         dataMock.fieldFormats,
         visData.columns[0],
-        colorIndexMap(ChartTypes.MOSAIC)
+        getColorIndexMap(ChartTypes.MOSAIC)
       );
       expect(registry.get().getCategoricalColor).toHaveBeenCalledWith(
         [expect.objectContaining({ name: 'Second level 1' })],

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

@@ -12,6 +12,7 @@ import { isEqual } from 'lodash';
 import type { PaletteRegistry, SeriesLayer, PaletteOutput, PaletteDefinition } from '@kbn/coloring';
 import type { FieldFormatsStart } from '@kbn/field-formats-plugin/public';
 import { lightenColor } from '@kbn/charts-plugin/public';
+import { SerializedValue } from '@kbn/data-plugin/common';
 import { BucketColumns, ChartTypes, PartitionVisParams } from '../../../common/types';
 import { DistinctSeries } from '../get_distinct_series';
 
@@ -123,7 +124,7 @@ const createSeriesLayers = (
   arrayNode: SimplifiedArrayNode,
   parentSeries: DistinctSeries['parentSeries'],
   isSplitChart: boolean,
-  colorIndexMap: Map<string, number>
+  colorIndexMap: Map<SerializedValue, number>
 ): SeriesLayer[] => {
   const seriesLayers: SeriesLayer[] = [];
   let tempParent: typeof arrayNode | (typeof arrayNode)['parent'] = arrayNode;
@@ -193,7 +194,7 @@ export const getColor = (
   isDarkMode: boolean,
   formatter: FieldFormatsStart,
   column: Partial<BucketColumns>,
-  colorIndexMap: Map<string, number>
+  colorIndexMap: Map<SerializedValue, number>
 ) => {
   // Mind the difference here: the contrast computation for the text ignores the alpha/opacity
   // therefore change it for dark mode

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

@@ -8,7 +8,7 @@
  */
 
 import { Datum, PartitionLayer } from '@elastic/charts';
-import { PaletteRegistry, getColorFactory } from '@kbn/coloring';
+import { ColorHandlingFn, PaletteRegistry, getColorFactory } from '@kbn/coloring';
 import { i18n } from '@kbn/i18n';
 import { FieldFormat } from '@kbn/field-formats-plugin/common';
 import type { FieldFormatsStart } from '@kbn/field-formats-plugin/public';
@@ -135,7 +135,7 @@ function getColorFromMappingFactory(
   palettes: KbnPalettes,
   visParams: PartitionVisParams,
   isDarkMode: boolean
-): undefined | ((category: string | string[]) => string) {
+): undefined | ColorHandlingFn {
   const { colorMapping, dimensions } = visParams;
 
   if (!colorMapping) {

src/platform/plugins/shared/chart_expressions/expression_tagcloud/public/components/tagcloud_component.tsx

@@ -23,12 +23,11 @@ import {
 } from '@elastic/charts';
 import { EmptyPlaceholder } from '@kbn/charts-plugin/public';
 import { useElasticChartsTheme } from '@kbn/charts-theme';
-import { PaletteRegistry, PaletteOutput, getColorFactory } from '@kbn/coloring';
+import { PaletteRegistry, PaletteOutput, getColorFactory, ColorHandlingFn } from '@kbn/coloring';
 import { IInterpreterRenderHandlers, DatatableRow } from '@kbn/expressions-plugin/public';
 import { getColorCategories, getOverridesFor } from '@kbn/chart-expressions-common';
 import type { AllowedSettingsOverrides, AllowedChartOverrides } from '@kbn/charts-plugin/common';
 import { getColumnByAccessor, getFormatByAccessor } from '@kbn/visualizations-plugin/common/utils';
-import { isMultiFieldKey } from '@kbn/data-plugin/common';
 import { KbnPalettes, useKbnPalettes } from '@kbn/palettes';
 import { css } from '@emotion/react';
 import { getFormatService } from '../format_service';
@@ -129,9 +128,11 @@ export const TagCloudChart = ({
     );
 
     return visData.rows.map((row) => {
-      const tag = tagColumn === undefined ? 'all' : row[tagColumn];
+      const { value: tagValue, tag } =
+        tagColumn === undefined
+          ? { value: undefined, tag: 'all' }
+          : { value: row[tagColumn], tag: row[tagColumn] };
 
-      const category = isMultiFieldKey(tag) ? tag.keys.map(String) : `${tag}`;
       return {
         text: bucketFormatter ? bucketFormatter.convert(tag, 'text') : tag,
         weight:
@@ -139,7 +140,7 @@ export const TagCloudChart = ({
             ? 1
             : calculateWeight(row[metricColumn], minValue, maxValue, 0, 1) || 0,
         color: colorFromMappingFn
-          ? colorFromMappingFn(category)
+          ? colorFromMappingFn(tagValue)
           : getColor(palettesRegistry, palette, tag, values, syncColors) || 'rgba(0,0,0,0)',
       };
     });
@@ -319,7 +320,7 @@ function getColorFromMappingFactory(
   palettes: KbnPalettes,
   isDarkMode: boolean,
   colorMapping?: string
-): undefined | ((category: string | string[]) => string) {
+): undefined | ColorHandlingFn {
   if (!colorMapping) {
     // return undefined, we will use the legacy color mapping instead
     return undefined;

src/platform/plugins/shared/chart_expressions/expression_tagcloud/tsconfig.json

@@ -27,7 +27,6 @@
     "@kbn/analytics",
     "@kbn/chart-expressions-common",
     "@kbn/chart-icons",
-    "@kbn/data-plugin",
     "@kbn/react-kibana-context-render",
     "@kbn/palettes",
     "@kbn/charts-theme",

src/platform/plugins/shared/chart_expressions/expression_xy/public/components/__snapshots__/xy_chart.test.tsx.snap

@@ -438,6 +438,15 @@ exports[`XYChart component it renders area 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -1534,6 +1543,15 @@ exports[`XYChart component it renders area 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -1978,6 +1996,15 @@ exports[`XYChart component it renders bar 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -3074,6 +3101,15 @@ exports[`XYChart component it renders bar 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -3518,6 +3554,15 @@ exports[`XYChart component it renders horizontal bar 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -4614,6 +4659,15 @@ exports[`XYChart component it renders horizontal bar 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -5058,6 +5112,15 @@ exports[`XYChart component it renders line 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -6154,6 +6217,15 @@ exports[`XYChart component it renders line 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -6598,6 +6670,15 @@ exports[`XYChart component it renders stacked area 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -7694,6 +7775,15 @@ exports[`XYChart component it renders stacked area 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -8138,6 +8228,15 @@ exports[`XYChart component it renders stacked bar 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -9234,6 +9333,15 @@ exports[`XYChart component it renders stacked bar 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -9678,6 +9786,15 @@ exports[`XYChart component it renders stacked horizontal bar 1`] = `
               "formattedColumns": Object {
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {},
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -10774,6 +10891,15 @@ exports[`XYChart component it renders stacked horizontal bar 1`] = `
                 "formattedColumns": Object {
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {},
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -11244,6 +11370,18 @@ exports[`XYChart component split chart should render split chart if both, splitR
                 "b": true,
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {
+                  2 => 2,
+                  5 => 5,
+                },
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -12540,6 +12678,18 @@ exports[`XYChart component split chart should render split chart if both, splitR
                   "b": true,
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {
+                    2 => 2,
+                    5 => 5,
+                  },
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -13024,6 +13174,18 @@ exports[`XYChart component split chart should render split chart if splitColumnA
                 "b": true,
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {
+                  2 => 2,
+                  5 => 5,
+                },
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -14313,6 +14475,18 @@ exports[`XYChart component split chart should render split chart if splitColumnA
                   "b": true,
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {
+                    2 => 2,
+                    5 => 5,
+                  },
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {
@@ -14795,6 +14969,18 @@ exports[`XYChart component split chart should render split chart if splitRowAcce
                 "b": true,
                 "c": true,
               },
+              "invertedRawValueMap": Map {
+                "a" => Map {},
+                "b" => Map {
+                  2 => 2,
+                  5 => 5,
+                },
+                "c" => Map {
+                  1652034840000 => 1652034840000,
+                  1652122440000 => 1652122440000,
+                },
+                "d" => Map {},
+              },
               "table": Object {
                 "columns": Array [
                   Object {
@@ -16084,6 +16270,18 @@ exports[`XYChart component split chart should render split chart if splitRowAcce
                   "b": true,
                   "c": true,
                 },
+                "invertedRawValueMap": Map {
+                  "a" => Map {},
+                  "b" => Map {
+                    2 => 2,
+                    5 => 5,
+                  },
+                  "c" => Map {
+                    1652034840000 => 1652034840000,
+                    1652122440000 => 1652122440000,
+                  },
+                  "d" => Map {},
+                },
                 "table": Object {
                   "columns": Array [
                     Object {

src/platform/plugins/shared/chart_expressions/expression_xy/public/components/data_layers.tsx

@@ -144,9 +144,6 @@ export const DataLayers: FC<Props> = ({
             ? JSON.parse(columnToLabel)
             : {};
 
-          // what if row values are not primitive? That is the case of, for instance, Ranges
-          // remaps them to their serialized version with the formatHint metadata
-          // In order to do it we need to make a copy of the table as the raw one is required for more features (filters, etc...) later on
           const formattedDatatableInfo = formattedDatatables[layerId];
 
           const yAxis = yAxesConfiguration.find((axisConfiguration) =>

src/platform/plugins/shared/chart_expressions/expression_xy/public/components/legend_action.test.tsx

@@ -19,7 +19,8 @@ import { getLegendAction } from './legend_action';
 import { LegendActionPopover, LegendCellValueActions } from './legend_action_popover';
 import { mockPaletteOutput } from '../../common/__mocks__';
 import { FieldFormat } from '@kbn/field-formats-plugin/common';
-import { LayerFieldFormats } from '../helpers';
+import { InvertedRawValueMap, LayerFieldFormats } from '../helpers';
+import { RawValue } from '@kbn/data-plugin/common';
 
 const legendCellValueActions: LegendCellValueActions = [
   { id: 'action_1', displayName: 'Action 1', iconType: 'testIcon1', execute: () => {} },
@@ -180,6 +181,9 @@ const sampleLayer: DataLayerConfig = {
 
 describe('getLegendAction', function () {
   let wrapperProps: LegendActionProps;
+  const invertedRawValueMap: InvertedRawValueMap = new Map(
+    table.columns.map((c) => [c.id, new Map<string, RawValue>()])
+  );
   const Component: React.ComponentType<LegendActionProps> = getLegendAction(
     [sampleLayer],
     jest.fn(),
@@ -201,6 +205,7 @@ describe('getLegendAction', function () {
     {
       first: {
         table,
+        invertedRawValueMap,
         formattedColumns: {},
       },
     },

src/platform/plugins/shared/chart_expressions/expression_xy/public/components/tooltip/tooltip.test.tsx

@@ -10,10 +10,16 @@
 import React from 'react';
 import { shallow } from 'enzyme';
 import { Tooltip } from './tooltip';
-import { generateSeriesId, LayersAccessorsTitles, LayersFieldFormats } from '../../helpers';
+import {
+  generateSeriesId,
+  InvertedRawValueMap,
+  LayersAccessorsTitles,
+  LayersFieldFormats,
+} from '../../helpers';
 import { XYChartSeriesIdentifier } from '@elastic/charts';
 import { sampleArgs, sampleLayer } from '../../../common/__mocks__';
 import { FieldFormat, FormatFactory } from '@kbn/field-formats-plugin/common';
+import { RawValue } from '@kbn/data-plugin/common';
 
 const getSeriesIdentifier = ({
   layerId,
@@ -44,6 +50,9 @@ const getSeriesIdentifier = ({
 
 describe('Tooltip', () => {
   const { data } = sampleArgs();
+  const invertedRawValueMap: InvertedRawValueMap = new Map(
+    data.columns.map((c) => [c.id, new Map<string, RawValue>()])
+  );
   const { layerId, xAccessor, splitAccessors = [], accessors } = sampleLayer;
   const seriesSplitAccessors = new Map();
   splitAccessors.forEach((splitAccessor) => {
@@ -112,7 +121,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         layers={[sampleLayer]}
       />
@@ -132,7 +143,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         xDomain={xDomain}
         layers={[sampleLayer]}
@@ -153,7 +166,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         xDomain={xDomain}
         layers={[sampleLayer]}
@@ -171,7 +186,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         xDomain={xDomain2}
         layers={[sampleLayer]}
@@ -191,7 +208,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         layers={[sampleLayer]}
       />
@@ -217,7 +236,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         layers={[sampleLayer]}
       />
@@ -245,7 +266,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         layers={[sampleLayer]}
       />
@@ -274,7 +297,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor, splitRowAccessor }}
         layers={[sampleLayer]}
       />
@@ -302,7 +327,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitColumnAccessor }}
         layers={[sampleLayer]}
       />
@@ -330,7 +357,9 @@ describe('Tooltip', () => {
         fieldFormats={fieldFormats}
         titles={titles}
         formatFactory={formatFactory}
-        formattedDatatables={{ [layerId]: { table: data, formattedColumns: {} } }}
+        formattedDatatables={{
+          [layerId]: { table: data, invertedRawValueMap, formattedColumns: {} },
+        }}
         splitAccessors={{ splitRowAccessor }}
         layers={[sampleLayer]}
       />

src/platform/plugins/shared/chart_expressions/expression_xy/public/helpers/color/color_mapping_accessor.ts

@@ -9,43 +9,32 @@
 
 import { SeriesColorAccessorFn } from '@elastic/charts';
 import { getColorFactory, type ColorMapping, type ColorMappingInputData } from '@kbn/coloring';
-import { MULTI_FIELD_KEY_SEPARATOR } from '@kbn/data-plugin/common';
 import { KbnPalettes } from '@kbn/palettes';
+import { InvertedRawValueMap } from '../data_layers';
 
 /**
  * Return a color accessor function for XY charts depending on the split accessors received.
  */
 export function getColorSeriesAccessorFn(
   config: ColorMapping.Config,
+  invertedRawValueMap: InvertedRawValueMap,
   palettes: KbnPalettes,
   isDarkMode: boolean,
   mappingData: ColorMappingInputData,
-  fieldId: string,
-  specialTokens: Map<string, string>
+  fieldId: string
 ): SeriesColorAccessorFn {
-  // inverse map to handle the conversion between the formatted string and their original format
-  // for any specified special tokens
-  const specialHandlingInverseMap: Map<string, string> = new Map(
-    [...specialTokens.entries()].map((d) => [d[1], d[0]])
-  );
-
   const getColor = getColorFactory(config, palettes, isDarkMode, mappingData);
+  const rawValueMap = invertedRawValueMap.get(fieldId) ?? new Map<string, unknown>();
 
   return ({ splitAccessors }) => {
     const splitValue = splitAccessors.get(fieldId);
-    // if there isn't a category associated in the split accessor, let's use the default color
-    if (splitValue === undefined) {
-      return null;
-    }
 
-    // category can be also a number, range, ip, multi-field. We need to stringify it to be sure
-    // we can correctly match it a with user string
-    // if the separator exist, we de-construct it into a multifieldkey into values.
-    const categories = `${splitValue}`.split(MULTI_FIELD_KEY_SEPARATOR).map((category) => {
-      return specialHandlingInverseMap.get(category) ?? category;
-    });
-    // we must keep the array nature of a multi-field key or just use a single string
-    // This is required because the rule stored are checked differently for single values or multi-values
-    return getColor(categories.length > 1 ? categories : categories[0]);
+    // No category associated in the split accessor, use the default color
+    if (splitValue === undefined) return null;
+
+    const rawValue =
+      typeof splitValue === 'string' ? rawValueMap.get(splitValue) ?? splitValue : splitValue;
+
+    return getColor(rawValue);
   };
 }

src/platform/plugins/shared/chart_expressions/expression_xy/public/helpers/color_assignment.test.ts

@@ -13,6 +13,7 @@ import { LayerTypes } from '../../common/constants';
 import { Datatable } from '@kbn/expressions-plugin/common';
 import { FieldFormat } from '@kbn/field-formats-plugin/common';
 import { LayersFieldFormats } from './layers';
+import { DatatablesWithFormatInfo } from './data_layers';
 
 describe('color_assignment', () => {
   const tables: Record<string, Datatable> = {
@@ -108,14 +109,16 @@ describe('color_assignment', () => {
     },
   } as unknown as LayersFieldFormats;
 
-  const formattedDatatables = {
+  const formattedDatatables: DatatablesWithFormatInfo = {
     first: {
       table: tables['1'],
       formattedColumns: {},
+      invertedRawValueMap: new Map(tables['1'].columns.map((c) => [c.id, new Map()])),
     },
     second: {
       table: tables['2'],
       formattedColumns: {},
+      invertedRawValueMap: new Map(tables['2'].columns.map((c) => [c.id, new Map()])),
     },
   };
 
@@ -184,6 +187,9 @@ describe('color_assignment', () => {
       const newFormattedDatatables = {
         first: {
           formattedColumns: formattedDatatables.first.formattedColumns,
+          invertedRawValueMap: new Map(
+            formattedDatatables.first.table.columns.map((c) => [c.id, new Map()])
+          ),
           table: {
             ...formattedDatatables.first.table,
             rows: [{ split1: complexObject }, { split1: 'abc' }],
@@ -212,6 +218,7 @@ describe('color_assignment', () => {
         first: {
           formattedColumns: formattedDatatables.first.formattedColumns,
           table: { ...formattedDatatables.first.table, columns: [] },
+          invertedRawValueMap: new Map(),
         },
         second: formattedDatatables.second,
       };
@@ -274,6 +281,9 @@ describe('color_assignment', () => {
             ...formattedDatatables.first.table,
             rows: [{ split1: 'abc' }, { split1: { aProp: 123 } }],
           },
+          invertedRawValueMap: new Map(
+            formattedDatatables.first.table.columns.map((c) => [c.id, new Map()])
+          ),
         },
         second: formattedDatatables.second,
       };
@@ -292,10 +302,11 @@ describe('color_assignment', () => {
 
     it('should handle missing columns', () => {
       const newLayers = [{ ...layers[0], table: { ...tables['1'], columns: [] } }, layers[1]];
-      const newFormattedDatatables = {
+      const newFormattedDatatables: DatatablesWithFormatInfo = {
         first: {
           formattedColumns: formattedDatatables.first.formattedColumns,
           table: { ...formattedDatatables.first.table, columns: [] },
+          invertedRawValueMap: new Map(),
         },
         second: formattedDatatables.second,
       };

src/platform/plugins/shared/chart_expressions/expression_xy/public/helpers/data_layers.tsx

@@ -25,9 +25,9 @@ import { Datatable } from '@kbn/expressions-plugin/common';
 import { getAccessorByDimension } from '@kbn/visualizations-plugin/common/utils';
 import type { ExpressionValueVisDimension } from '@kbn/visualizations-plugin/common/expression_functions';
 import { PaletteRegistry, SeriesLayer } from '@kbn/coloring';
-import { SPECIAL_TOKENS_STRING_CONVERSION } from '@kbn/coloring';
 import { getColorCategories } from '@kbn/chart-expressions-common';
 import { KbnPalettes } from '@kbn/palettes';
+import { RawValue } from '@kbn/data-plugin/common';
 import { isDataLayer } from '../../common/utils/layer_types_guards';
 import { CommonXYDataLayerConfig, CommonXYLayerConfig, XScaleType } from '../../common';
 import { AxisModes, SeriesTypes } from '../../common/constants';
@@ -40,6 +40,7 @@ import { getFormat } from './format';
 import { getColorSeriesAccessorFn } from './color/color_mapping_accessor';
 
 type SeriesSpec = LineSeriesProps & BarSeriesProps & AreaSeriesProps;
+export type InvertedRawValueMap = Map<string, Map<string, RawValue>>;
 
 type GetSeriesPropsFn = (config: {
   layer: CommonXYDataLayerConfig;
@@ -109,6 +110,10 @@ type GetLineConfigFn = (config: {
 export interface DatatableWithFormatInfo {
   table: Datatable;
   formattedColumns: Record<string, true>;
+  /**
+   * Inverse map per column to link formatted string to complex values (i.e. `RangeKey`).
+   */
+  invertedRawValueMap: InvertedRawValueMap;
 }
 
 export type DatatablesWithFormatInfo = Record<string, DatatableWithFormatInfo>;
@@ -124,7 +129,8 @@ export const getFormattedRow = (
   xAccessor: string | undefined,
   splitColumnAccessor: string | undefined,
   splitRowAccessor: string | undefined,
-  xScaleType: XScaleType
+  xScaleType: XScaleType,
+  invertedRawValueMap: InvertedRawValueMap
 ): { row: Datatable['rows'][number]; formattedColumns: Record<string, true> } =>
   columns.reduce(
     (formattedInfo, { id }) => {
@@ -137,8 +143,10 @@ export const getFormattedRow = (
           id === splitColumnAccessor ||
           id === splitRowAccessor)
       ) {
+        const formattedValue = columnsFormatters[id]?.convert(record) ?? '';
+        invertedRawValueMap.get(id)?.set(formattedValue, record);
         return {
-          row: { ...formattedInfo.row, [id]: columnsFormatters[id]!.convert(record) },
+          row: { ...formattedInfo.row, [id]: formattedValue },
           formattedColumns: { ...formattedInfo.formattedColumns, [id]: true },
         };
       }
@@ -155,7 +163,7 @@ export const getFormattedTable = (
   splitRowAccessor: string | ExpressionValueVisDimension | undefined,
   accessors: Array<string | ExpressionValueVisDimension>,
   xScaleType: XScaleType
-): { table: Datatable; formattedColumns: Record<string, true> } => {
+): DatatableWithFormatInfo => {
   const columnsFormatters = table.columns.reduce<Record<string, IFieldFormat>>(
     (formatters, { id, meta }) => {
       const accessor: string | ExpressionValueVisDimension | undefined = accessors.find(
@@ -170,6 +178,9 @@ export const getFormattedTable = (
     {}
   );
 
+  const invertedRawValueMap: InvertedRawValueMap = new Map(
+    table.columns.map((c) => [c.id, new Map<string, RawValue>()])
+  );
   const formattedTableInfo: {
     rows: Datatable['rows'];
     formattedColumns: Record<string, true>;
@@ -185,7 +196,8 @@ export const getFormattedTable = (
       xAccessor ? getAccessorByDimension(xAccessor, table.columns) : undefined,
       splitColumnAccessor ? getAccessorByDimension(splitColumnAccessor, table.columns) : undefined,
       splitRowAccessor ? getAccessorByDimension(splitRowAccessor, table.columns) : undefined,
-      xScaleType
+      xScaleType,
+      invertedRawValueMap
     );
     formattedTableInfo.rows.push(formattedRowInfo.row);
     formattedTableInfo.formattedColumns = {
@@ -195,6 +207,7 @@ export const getFormattedTable = (
   }
 
   return {
+    invertedRawValueMap,
     table: { ...table, rows: formattedTableInfo.rows },
     formattedColumns: formattedTableInfo.formattedColumns,
   };
@@ -438,10 +451,7 @@ export const getSeriesProps: GetSeriesPropsFn = ({
     markSizeAccessor ? getFormat(table.columns, markSizeAccessor) : undefined
   );
 
-  // what if row values are not primitive? That is the case of, for instance, Ranges
-  // remaps them to their serialized version with the formatHint metadata
-  // In order to do it we need to make a copy of the table as the raw one is required for more features (filters, etc...) later on
-  const { table: formattedTable, formattedColumns } = formattedDatatableInfo;
+  const { table: formattedTable, formattedColumns, invertedRawValueMap } = formattedDatatableInfo;
 
   // For date histogram chart type, we're getting the rows that represent intervals without data.
   // To not display them in the legend, they need to be filtered out.
@@ -488,14 +498,14 @@ export const getSeriesProps: GetSeriesPropsFn = ({
     layer.colorMapping && splitColumnIds.length > 0
       ? getColorSeriesAccessorFn(
           JSON.parse(layer.colorMapping), // the color mapping is at this point just a stringified JSON
+          invertedRawValueMap,
           palettes,
           isDarkMode,
           {
             type: 'categories',
             categories: getColorCategories(table.rows, splitColumnIds[0]),
           },
-          splitColumnIds[0],
-          SPECIAL_TOKENS_STRING_CONVERSION
+          splitColumnIds[0]
         )
       : (series) =>
           getColor(

src/platform/plugins/shared/data/common/index.ts

@@ -7,8 +7,6 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-// TODO: https://github.com/elastic/kibana/issues/109904
-
 export {
   DEFAULT_QUERY_LANGUAGE,
   KIBANA_USER_QUERY_LANGUAGE_KEY,
@@ -17,6 +15,10 @@ export {
   SCRIPT_LANGUAGES_ROUTE_LATEST_VERSION,
   UI_SETTINGS,
 } from './constants';
+export type { RawValue } from './serializable_field';
+export { SerializableField } from './serializable_field';
+export type { SerializedField, SerializedValue } from './serialize_utils';
+export { SerializableType, deserializeField, serializeField } from './serialize_utils';
 export type { ValueSuggestionsMethod } from './constants';
 export { DatatableUtilitiesService } from './datatable_utilities';
 export { getEsQueryConfig } from './es_query';
@@ -309,7 +311,7 @@ export {
   termsAggFilter,
   getTermsBucketAgg,
   MultiFieldKey,
-  isMultiFieldKey,
+  RangeKey,
   MULTI_FIELD_KEY_SEPARATOR,
   aggMultiTermsFnName,
   aggMultiTerms,

src/platform/plugins/shared/data/common/search/aggs/buckets/__snapshots__/range.test.ts.snap

@@ -0,0 +1,17 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`Range Agg RangeKey fully closed range should correctly stringify field 1`] = `"from:0,to:100"`;
+
+exports[`Range Agg RangeKey fully closed range w/ label should correctly stringify field 1`] = `"from:0,to:100"`;
+
+exports[`Range Agg RangeKey open lower range should correctly stringify field 1`] = `"from:-Infinity,to:100"`;
+
+exports[`Range Agg RangeKey open lower range w/ label should correctly stringify field 1`] = `"from:-Infinity,to:100"`;
+
+exports[`Range Agg RangeKey open range should correctly stringify field 1`] = `"from:-Infinity,to:Infinity"`;
+
+exports[`Range Agg RangeKey open range should correctly stringify field 2`] = `"from:-Infinity,to:Infinity"`;
+
+exports[`Range Agg RangeKey open upper range should correctly stringify field 1`] = `"from:0,to:Infinity"`;
+
+exports[`Range Agg RangeKey open upper range w/ label should correctly stringify field 1`] = `"from:0,to:Infinity"`;

src/platform/plugins/shared/data/common/search/aggs/buckets/index.ts

@@ -34,13 +34,14 @@ export { TimeBuckets, convertDurationToNormalizedEsInterval } from './lib/time_b
 export * from './migrate_include_exclude_format';
 export * from './range_fn';
 export * from './range';
+export * from './range_key';
 export * from './significant_terms_fn';
 export * from './significant_terms';
 export * from './significant_text_fn';
 export * from './significant_text';
 export * from './terms_fn';
 export * from './terms';
-export { MultiFieldKey, isMultiFieldKey, MULTI_FIELD_KEY_SEPARATOR } from './multi_field_key';
+export { MultiFieldKey, MULTI_FIELD_KEY_SEPARATOR } from './multi_field_key';
 export * from './multi_terms_fn';
 export * from './multi_terms';
 export * from './rare_terms_fn';

src/platform/plugins/shared/data/common/search/aggs/buckets/multi_field_key.ts

@@ -7,7 +7,16 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-const id = Symbol('id');
+import { SerializableField } from '../../../serializable_field';
+import { SerializableType } from '../../../serialize_utils';
+
+/**
+ * Serialized form of {@link @kbn/data-plugin/common.MultiFieldKey}
+ */
+export interface SerializedMultiFieldKey {
+  type: typeof SerializableType.MultiFieldKey;
+  keys: string[];
+}
 
 const isBucketLike = (bucket: unknown): bucket is { key: unknown } => {
   return Boolean(bucket && typeof bucket === 'object' && 'key' in bucket);
@@ -17,31 +26,41 @@ function getKeysFromBucket(bucket: unknown) {
   if (!isBucketLike(bucket)) {
     throw new Error('bucket malformed - no key found');
   }
-  return Array.isArray(bucket.key)
-    ? bucket.key.map((keyPart) => String(keyPart))
-    : [String(bucket.key)];
+  return Array.isArray(bucket.key) ? bucket.key.map(String) : [String(bucket.key)];
 }
 
-export class MultiFieldKey {
-  [id]: string;
-  keys: string[];
-
-  constructor(bucket: unknown) {
-    this.keys = getKeysFromBucket(bucket);
-
-    this[id] = MultiFieldKey.idBucket(bucket);
+export class MultiFieldKey extends SerializableField<SerializedMultiFieldKey> {
+  static isInstance(field: unknown): field is MultiFieldKey {
+    return field instanceof MultiFieldKey;
   }
-  static idBucket(bucket: unknown) {
+
+  static deserialize(value: SerializedMultiFieldKey): MultiFieldKey {
+    return new MultiFieldKey({
+      key: value.keys, // key here is to keep bwc with constructor params
+    });
+  }
+
+  static idBucket(bucket: unknown): string {
     return getKeysFromBucket(bucket).join(',');
   }
 
-  toString() {
-    return this[id];
-  }
-}
+  keys: string[];
 
-export function isMultiFieldKey(field: unknown): field is MultiFieldKey {
-  return field instanceof MultiFieldKey;
+  constructor(bucket: unknown) {
+    super();
+    this.keys = getKeysFromBucket(bucket);
+  }
+
+  toString(): string {
+    return this.keys.join(',');
+  }
+
+  serialize(): SerializedMultiFieldKey {
+    return {
+      type: SerializableType.MultiFieldKey,
+      keys: this.keys,
+    };
+  }
 }
 
 /**

src/platform/plugins/shared/data/common/search/aggs/buckets/range.test.ts

@@ -11,55 +11,108 @@ import { AggConfigs } from '../agg_configs';
 import { mockAggTypesRegistry } from '../test_helpers';
 import { BUCKET_TYPES } from './bucket_agg_types';
 import { FieldFormatsGetConfigFn, NumberFormat } from '@kbn/field-formats-plugin/common';
+import { RangeKey } from './range_key';
 
 describe('Range Agg', () => {
-  const getConfig = (() => {}) as FieldFormatsGetConfigFn;
-  const getAggConfigs = () => {
-    const field = {
-      name: 'bytes',
+  describe('RangeKey', () => {
+    const label = 'some label';
+    describe.each([
+      ['open range', {}, undefined],
+      ['open upper range', { from: 0 }, undefined],
+      ['open lower range', { to: 100 }, undefined],
+      ['fully closed range', { from: 0, to: 100 }, undefined],
+      ['open range', {}, [{ label }]],
+      ['open upper range w/ label', { from: 0 }, [{ from: 0, label }]],
+      ['open lower range w/ label', { to: 100 }, [{ to: 100, label }]],
+      ['fully closed range w/ label', { from: 0, to: 100 }, [{ from: 0, to: 100, label }]],
+    ])('%s', (_, bucket: any, ranges) => {
+      const initial = new RangeKey(bucket, ranges);
+
+      test('should correctly set gte', () => {
+        expect(initial.gte).toBe(bucket?.from == null ? -Infinity : bucket.from);
+      });
+
+      test('should correctly set lt', () => {
+        expect(initial.lt).toBe(bucket?.to == null ? Infinity : bucket.to);
+      });
+
+      test('should correctly set label', () => {
+        expect(initial.label).toBe(ranges?.[0]?.label);
+      });
+
+      test('should correctly stringify field', () => {
+        expect(initial.toString()).toMatchSnapshot();
+      });
+    });
+  });
+
+  describe('#fromString', () => {
+    test.each([
+      ['empty range', '', {}],
+      ['bad buckets', 'from:baddd,to:baddd', {}],
+      ['open range', 'from:undefined,to:undefined', {}],
+      ['open upper range', 'from:0,to:undefined', { from: 0 }],
+      ['open lower range', 'from:undefined,to:100', { to: 100 }],
+      ['fully closed range', 'from:0,to:100', { from: 0, to: 100 }],
+      ['mixed closed range', 'from:-100,to:100', { from: -100, to: 100 }],
+      ['mixed open range', 'from:-100,to:undefined', { from: -100 }],
+      ['negative closed range', 'from:-100,to:-50', { from: -100, to: -50 }],
+      ['negative open range', 'from:undefined,to:-50', { to: -50 }],
+    ])('should correctly build RangeKey from string for %s', (_, rangeString, bucket) => {
+      const expected = new RangeKey(bucket);
+
+      expect(RangeKey.fromString(rangeString).toString()).toBe(expected.toString());
+    });
+  });
+
+  describe('RangeKey with getAggConfigs', () => {
+    const getConfig = (() => {}) as FieldFormatsGetConfigFn;
+    const getAggConfigs = () => {
+      const field = {
+        name: 'bytes',
+      };
+
+      const indexPattern = {
+        id: '1234',
+        title: 'logstash-*',
+        fields: {
+          getByName: () => field,
+          filter: () => [field],
+        },
+        getFormatterForField: () =>
+          new NumberFormat(
+            {
+              pattern: '0,0.[000] b',
+            },
+            getConfig
+          ),
+      } as any;
+
+      return new AggConfigs(
+        indexPattern,
+        [
+          {
+            type: BUCKET_TYPES.RANGE,
+            schema: 'segment',
+            params: {
+              field: 'bytes',
+              ranges: [
+                { from: 0, to: 1000 },
+                { from: 1000, to: 2000 },
+              ],
+            },
+          },
+        ],
+        {
+          typesRegistry: mockAggTypesRegistry(),
+        },
+        jest.fn()
+      );
     };
 
-    const indexPattern = {
-      id: '1234',
-      title: 'logstash-*',
-      fields: {
-        getByName: () => field,
-        filter: () => [field],
-      },
-      getFormatterForField: () =>
-        new NumberFormat(
-          {
-            pattern: '0,0.[000] b',
-          },
-          getConfig
-        ),
-    } as any;
-
-    return new AggConfigs(
-      indexPattern,
-      [
-        {
-          type: BUCKET_TYPES.RANGE,
-          schema: 'segment',
-          params: {
-            field: 'bytes',
-            ranges: [
-              { from: 0, to: 1000 },
-              { from: 1000, to: 2000 },
-            ],
-          },
-        },
-      ],
-      {
-        typesRegistry: mockAggTypesRegistry(),
-      },
-      jest.fn()
-    );
-  };
-
-  test('produces the expected expression ast', () => {
-    const aggConfigs = getAggConfigs();
-    expect(aggConfigs.aggs[0].toExpressionAst()).toMatchInlineSnapshot(`
+    test('produces the expected expression ast', () => {
+      const aggConfigs = getAggConfigs();
+      expect(aggConfigs.aggs[0].toExpressionAst()).toMatchInlineSnapshot(`
       Object {
         "chain": Array [
           Object {
@@ -120,23 +173,24 @@ describe('Range Agg', () => {
         "type": "expression",
       }
     `);
-  });
+    });
 
-  describe('getSerializedFormat', () => {
-    test('generates a serialized field format in the expected shape', () => {
-      const aggConfigs = getAggConfigs();
-      const agg = aggConfigs.aggs[0];
-      expect(agg.type.getSerializedFormat(agg)).toMatchInlineSnapshot(`
-        Object {
-          "id": "range",
-          "params": Object {
-            "id": "number",
+    describe('getSerializedFormat', () => {
+      test('generates a serialized field format in the expected shape', () => {
+        const aggConfigs = getAggConfigs();
+        const agg = aggConfigs.aggs[0];
+        expect(agg.type.getSerializedFormat(agg)).toMatchInlineSnapshot(`
+          Object {
+            "id": "range",
             "params": Object {
-              "pattern": "0,0.[000] b",
+              "id": "number",
+              "params": Object {
+                "pattern": "0,0.[000] b",
+              },
             },
-          },
-        }
-      `);
+          }
+        `);
+      });
     });
   });
 });

src/platform/plugins/shared/data/common/search/aggs/buckets/range_key.ts

@@ -7,47 +7,121 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-const id = Symbol('id');
+import { SerializableField } from '../../../serializable_field';
+import { SerializableType } from '../../../serialize_utils';
 
 type Ranges = Array<
   Partial<{
-    from: string | number;
-    to: string | number;
+    from: string | number | null;
+    to: string | number | null;
     label: string;
   }>
 >;
 
-export class RangeKey {
-  [id]: string;
+type RangeValue = string | number | undefined | null;
+interface BucketLike {
+  from?: RangeValue;
+  to?: RangeValue;
+}
+
+/**
+ * Serialized form of {@link @kbn/data-plugin/common.RangeKey}
+ */
+export interface SerializedRangeKey {
+  type: typeof SerializableType.RangeKey;
+  from: string | number | null;
+  to: string | number | null;
+  ranges: Ranges;
+}
+
+function findCustomLabel(from: RangeValue, to: RangeValue, ranges?: Ranges) {
+  return (ranges || []).find(
+    (range) =>
+      ((from == null && range.from == null) || range.from === from) &&
+      ((to == null && range.to == null) || range.to === to)
+  )?.label;
+}
+
+const getRangeValue = (bucket: unknown, key: string): RangeValue => {
+  const value = bucket && typeof bucket === 'object' && key in bucket && (bucket as any)[key];
+  return value == null || ['string', 'number'].includes(typeof value) ? value : null;
+};
+
+const getRangeFromBucket = (bucket: unknown): BucketLike => {
+  return {
+    from: getRangeValue(bucket, 'from'),
+    to: getRangeValue(bucket, 'to'),
+  };
+};
+
+const regex = /^from:(-?\d+?|undefined),to:(-?\d+?|undefined)$/;
+
+export class RangeKey extends SerializableField<SerializedRangeKey> {
+  static isInstance(field: unknown): field is RangeKey {
+    return field instanceof RangeKey;
+  }
+
+  static deserialize(value: SerializedRangeKey): RangeKey {
+    const { to, from, ranges } = value;
+    return new RangeKey({ to, from }, ranges);
+  }
+
+  static idBucket(bucket: unknown): string {
+    const { from, to } = getRangeFromBucket(bucket);
+    return `from:${from},to:${to}`;
+  }
+
+  static isRangeKeyString(rangeKey: string): boolean {
+    return regex.test(rangeKey);
+  }
+
+  /**
+   * Returns `RangeKey` from stringified form. Cannot extract labels from stringified form.
+   *
+   * Only supports numerical (non-string) values.
+   */
+  static fromString(rangeKey: string): RangeKey {
+    const [from, to] = (regex.exec(rangeKey) ?? [])
+      .slice(1)
+      .map(Number)
+      .map((n) => (isNaN(n) ? undefined : n));
+
+    return new RangeKey({ from, to });
+  }
+
   gte: string | number;
   lt: string | number;
   label?: string;
 
-  private findCustomLabel(
-    from: string | number | undefined | null,
-    to: string | number | undefined | null,
-    ranges?: Ranges
-  ) {
-    return (ranges || []).find(
-      (range) =>
-        ((from == null && range.from == null) || range.from === from) &&
-        ((to == null && range.to == null) || range.to === to)
-    )?.label;
+  constructor(bucket: unknown, allRanges?: Ranges) {
+    super();
+    const { from, to } = getRangeFromBucket(bucket);
+    this.gte = from == null ? -Infinity : from;
+    this.lt = to == null ? +Infinity : to;
+    this.label = findCustomLabel(from, to, allRanges);
   }
 
-  constructor(bucket: any, allRanges?: Ranges) {
-    this.gte = bucket.from == null ? -Infinity : bucket.from;
-    this.lt = bucket.to == null ? +Infinity : bucket.to;
-    this.label = this.findCustomLabel(bucket.from, bucket.to, allRanges);
-
-    this[id] = RangeKey.idBucket(bucket);
+  toString(): string {
+    return `from:${this.gte},to:${this.lt}`;
   }
 
-  static idBucket(bucket: any) {
-    return `from:${bucket.from},to:${bucket.to}`;
-  }
-
-  toString() {
-    return this[id];
+  serialize(): SerializedRangeKey {
+    const from = typeof this.gte === 'string' || isFinite(this.gte) ? this.gte : null;
+    const to = typeof this.lt === 'string' || isFinite(this.lt) ? this.lt : null;
+    return {
+      type: SerializableType.RangeKey,
+      from,
+      to,
+      ranges:
+        this.label === undefined
+          ? []
+          : [
+              {
+                from,
+                to,
+                label: this.label,
+              },
+            ],
+    };
   }
 }

src/platform/plugins/shared/data/common/serializable_field.ts

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+/**
+ * Alias for unknown raw field value, could be instance of a field Class
+ */
+export type RawValue = number | string | unknown;
+
+/**
+ * Class to extends that enabled serializing and deserializing instance values
+ */
+export abstract class SerializableField<S> {
+  static isSerializable<T>(field: RawValue): field is SerializableField<T> {
+    return Boolean((field as SerializableField<T>).serialize);
+  }
+
+  /**
+   * Serializes the class instance to a known `SerializedValue` that can be used to instantiate a new instance
+   *
+   * Ideally this returns the same params as found in the constructor.
+   */
+  abstract serialize(): S;
+
+  /**
+   * typescript forbids abstract static methods but this is a workaround to require it
+   *
+   * @param serializedValue type of `SerializedValue`
+   * @returns `instanceValue` should same type as instantiating class
+   */
+  static deserialize(serializedValue: unknown): unknown {
+    throw new Error(
+      'Must implement a static `deserialize` method to conform to /`SerializableField/`'
+    );
+  }
+}

src/platform/plugins/shared/data/common/serialize_utils.test.ts

@@ -0,0 +1,74 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { MultiFieldKey, RangeKey } from '.';
+import { SerializedRangeKey } from './search';
+import { SerializedMultiFieldKey } from './search/aggs/buckets/multi_field_key';
+import { SerializableType, deserializeField, serializeField } from './serialize_utils';
+
+describe('serializeField/deserializeField', () => {
+  describe('MultiFieldKey', () => {
+    it.each([
+      ['single value', { key: 'one' }],
+      ['multiple values', { key: ['one', 'two', 'three'] }],
+    ])('should serialize and deserialize %s', (_, bucket) => {
+      const initial = new MultiFieldKey(bucket);
+      const serialized = serializeField(initial) as SerializedMultiFieldKey;
+      expect(serialized.type).toBe(SerializableType.MultiFieldKey);
+      const deserialized = deserializeField(serialized) as MultiFieldKey;
+      expect(deserialized).toMatchObject(initial);
+      expect(deserialized.toString()).toBe(initial.toString());
+      expect(deserialized).toBeInstanceOf(MultiFieldKey);
+    });
+  });
+
+  describe('RangeKey', () => {
+    const label = 'some label';
+    it.each([
+      ['open range', {}, undefined],
+      ['open upper range', { from: 0 }, undefined],
+      ['open lower range', { to: 100 }, undefined],
+      ['fully closed range', { from: 0, to: 100 }, undefined],
+      ['open range', {}, [{ label }]],
+      ['open upper range w/ label', { from: 0 }, [{ from: 0, label }]],
+      ['open lower range w/ label', { to: 100 }, [{ to: 100, label }]],
+      ['fully closed range w/ label', { from: 0, to: 100 }, [{ from: 0, to: 100, label }]],
+    ])('should serialize and deserialize %s', (_, bucket, ranges) => {
+      const initial = new RangeKey(bucket, ranges);
+      const serialized = serializeField(initial) as SerializedRangeKey;
+      expect(serialized.type).toBe(SerializableType.RangeKey);
+      expect(serialized.ranges).toHaveLength(initial.label ? 1 : 0);
+      const deserialized = deserializeField(serialized) as RangeKey;
+      expect(RangeKey.idBucket(deserialized)).toBe(RangeKey.idBucket(initial));
+      expect(deserialized.gte).toBe(initial.gte);
+      expect(deserialized.lt).toBe(initial.lt);
+      expect(deserialized.label).toBe(initial.label);
+      expect(deserialized).toBeInstanceOf(RangeKey);
+    });
+  });
+
+  describe('Primitive values', () => {
+    it.each([
+      ['strings', 'some string'],
+      ['strings (empty)', ''],
+      ['numbers', 123],
+      ['numbers (0)', 0],
+      ['boolean (true)', true],
+      ['boolean (false)', false],
+      ['object', { test: 1 }],
+      ['array', ['test', 1]],
+      ['undefined', undefined],
+      ['null', null],
+    ])('should deserialize %s', (_, initial) => {
+      const serialized = serializeField(initial);
+      const deserialized = deserializeField(serialized);
+      expect(deserialized).toEqual(initial);
+    });
+  });
+});

src/platform/plugins/shared/data/common/serialize_utils.ts

@@ -0,0 +1,55 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { MultiFieldKey, RangeKey, SerializedRangeKey } from './search';
+import { SerializedMultiFieldKey } from './search/aggs/buckets/multi_field_key';
+import { RawValue, SerializableField } from './serializable_field';
+
+/**
+ * All available serialized forms of complex/instance fields. Excludes non-complex/primitive fields.
+ *
+ * Use `SerializedValue` for all generalize serial values which includes non-complex/primitive fields.
+ *
+ * Currently includes:
+ * - `RangeKey`
+ * - `MultiFieldKey`
+ */
+export type SerializedField = SerializedMultiFieldKey | SerializedRangeKey;
+
+/**
+ * Alias for unknown serialized value. This value is what we store in the SO and app state
+ * to persist the color assignment based on the raw row value.
+ *
+ * In most cases this is a `string` or `number` or plain `object`, in other cases this is an
+ * object serialized from an instance of a given field (i.e. `RangeKey` or `MultiFieldKey`).
+ */
+export type SerializedValue = number | string | SerializedField | unknown;
+
+export const SerializableType = {
+  MultiFieldKey: 'multiFieldKey' as const,
+  RangeKey: 'rangeKey' as const,
+};
+
+export function deserializeField(field: SerializedValue) {
+  const type = field != null && (field as any)?.type;
+
+  switch (type) {
+    case SerializableType.MultiFieldKey:
+      return MultiFieldKey.deserialize(field as SerializedMultiFieldKey);
+    case SerializableType.RangeKey:
+      return RangeKey.deserialize(field as SerializedRangeKey);
+    default:
+      return field;
+  }
+}
+
+export function serializeField(field: RawValue): SerializedValue {
+  if (field == null || !SerializableField.isSerializable(field)) return field;
+  return field.serialize();
+}

x-pack/platform/plugins/private/translations/translations/fr-FR.json

@@ -335,7 +335,6 @@
     "coloring.colorMapping.colorPicker.hexColorinputAriaLabel": "Entrée de couleur hex",
     "coloring.colorMapping.colorPicker.invalidColorHex": "Veuillez utiliser un code de couleur hex valide",
     "coloring.colorMapping.colorPicker.lowContrastColor": "La couleur est peu contrastée dans {themes} {errorModes, plural, one {mode} other {# modes}}",
-    "coloring.colorMapping.colorPicker.newColorAriaLabel": "Sélectionner une nouvelle couleur",
     "coloring.colorMapping.colorPicker.paletteColorsLabel": "Couleurs de la palette",
     "coloring.colorMapping.colorPicker.paletteTabLabel": "Couleurs",
     "coloring.colorMapping.colorPicker.pickAColorAriaLabel": "Sélectionner une couleur",
@@ -24492,8 +24491,6 @@
     "xpack.lens.colorMapping.editColorMappingTitle": "Attribuer des couleurs aux termes",
     "xpack.lens.colorMapping.editColors": "Modifier les couleurs",
     "xpack.lens.colorMapping.editColorsTitle": "Modifier les couleurs",
-    "xpack.lens.colorMapping.techPreviewLabel": "Préversion technique",
-    "xpack.lens.colorMapping.tryLabel": "Utiliser la nouvelle fonctionnalité de mapping des couleurs",
     "xpack.lens.colorSiblingFlyoutTitle": "Couleur",
     "xpack.lens.config.applyFlyoutLabel": "Appliquer et fermer",
     "xpack.lens.config.cancelFlyoutAriaLabel": "Annuler les changements appliqués",

x-pack/platform/plugins/private/translations/translations/ja-JP.json

@@ -335,7 +335,6 @@
     "coloring.colorMapping.colorPicker.hexColorinputAriaLabel": "16進数の色入力",
     "coloring.colorMapping.colorPicker.invalidColorHex": "有効な16進数の色コードを使用してください",
     "coloring.colorMapping.colorPicker.lowContrastColor": "この色は{themes} {errorModes, plural, other {# モード}}ではコントラストが低くなります。",
-    "coloring.colorMapping.colorPicker.newColorAriaLabel": "新しい色を選択",
     "coloring.colorMapping.colorPicker.paletteColorsLabel": "パレットの色",
     "coloring.colorMapping.colorPicker.paletteTabLabel": "色",
     "coloring.colorMapping.colorPicker.pickAColorAriaLabel": "色を選択",
@@ -24472,8 +24471,6 @@
     "xpack.lens.colorMapping.editColorMappingTitle": "色を用語に割り当て",
     "xpack.lens.colorMapping.editColors": "色を編集",
     "xpack.lens.colorMapping.editColorsTitle": "色を編集",
-    "xpack.lens.colorMapping.techPreviewLabel": "テクニカルプレビュー",
-    "xpack.lens.colorMapping.tryLabel": "新しい色マッピング機能を使用",
     "xpack.lens.colorSiblingFlyoutTitle": "色",
     "xpack.lens.config.applyFlyoutLabel": "適用して閉じる",
     "xpack.lens.config.cancelFlyoutAriaLabel": "適用された変更をキャンセル",

x-pack/platform/plugins/private/translations/translations/zh-CN.json

@@ -334,7 +334,6 @@
     "coloring.colorMapping.colorPicker.hexColorinputAriaLabel": "HEX 颜色输入",
     "coloring.colorMapping.colorPicker.invalidColorHex": "请使用有效的颜色 HEX 代码",
     "coloring.colorMapping.colorPicker.lowContrastColor": "此颜色在 {themes} {errorModes, plural, other {# 个模式}}下的对比度较低",
-    "coloring.colorMapping.colorPicker.newColorAriaLabel": "选择新颜色",
     "coloring.colorMapping.colorPicker.paletteColorsLabel": "调色板颜色",
     "coloring.colorMapping.colorPicker.paletteTabLabel": "颜色",
     "coloring.colorMapping.colorPicker.pickAColorAriaLabel": "选取颜色",
@@ -24518,8 +24517,6 @@
     "xpack.lens.colorMapping.editColorMappingTitle": "为词分配颜色",
     "xpack.lens.colorMapping.editColors": "编辑颜色",
     "xpack.lens.colorMapping.editColorsTitle": "编辑颜色",
-    "xpack.lens.colorMapping.techPreviewLabel": "技术预览",
-    "xpack.lens.colorMapping.tryLabel": "使用新的颜色映射功能",
     "xpack.lens.colorSiblingFlyoutTitle": "颜色",
     "xpack.lens.config.applyFlyoutLabel": "应用并关闭",
     "xpack.lens.config.cancelFlyoutAriaLabel": "取消应用的更改",

x-pack/platform/plugins/shared/lens/common/expressions/impl/datatable/utils.ts

@@ -44,6 +44,18 @@ export function isNumericField(meta?: DatatableColumnMeta): boolean {
   );
 }
 
+export function getDatatableColumn(table: Datatable | undefined, accessor: string) {
+  return table?.columns.find((col) => col.id === accessor || getOriginalId(col.id) === accessor);
+}
+
+export function getFieldMetaFromDatatable(table: Datatable | undefined, accessor: string) {
+  return getDatatableColumn(table, accessor)?.meta;
+}
+
+export function getFieldTypeFromDatatable(table: Datatable | undefined, accessor: string) {
+  return getFieldMetaFromDatatable(table, accessor)?.type;
+}
+
 /**
  * Returns true for numerical fields, excluding ranges
  *
@@ -54,8 +66,3 @@ export function isNumericFieldForDatatable(table: Datatable | undefined, accesso
   const meta = getFieldMetaFromDatatable(table, accessor);
   return isNumericField(meta);
 }
-
-export function getFieldMetaFromDatatable(table: Datatable | undefined, accessor: string) {
-  return table?.columns.find((col) => col.id === accessor || getOriginalId(col.id) === accessor)
-    ?.meta;
-}

x-pack/platform/plugins/shared/lens/public/app_plugin/app.tsx

@@ -176,7 +176,6 @@ export function App({
       setIndicateNoData(false);
     }
   }, [setIndicateNoData, indicateNoData, searchSessionId]);
-
   const getIsByValueMode = useCallback(
     () => Boolean(isLinkedToOriginatingApp && !savedObjectId),
     [isLinkedToOriginatingApp, savedObjectId]

x-pack/platform/plugins/shared/lens/public/app_plugin/lens_document_equality.ts

@@ -63,12 +63,13 @@ export const isLensEqual = (
   const availableDatasourceTypes1 = Object.keys(doc1.state.datasourceStates);
   const availableDatasourceTypes2 = Object.keys(doc2.state.datasourceStates);
 
+  // simple comparison
   let datasourcesEqual =
     intersection(availableDatasourceTypes1, availableDatasourceTypes2).length ===
     union(availableDatasourceTypes1, availableDatasourceTypes2).length;
 
   if (datasourcesEqual) {
-    // equal so far, so actually check
+    // deep comparison
     datasourcesEqual = availableDatasourceTypes1.every((type) =>
       datasourceMap[type].isEqual(
         doc1.state.datasourceStates[type],

x-pack/platform/plugins/shared/lens/public/app_plugin/shared/edit_on_the_fly/lens_configuration_flyout.tsx

@@ -119,14 +119,14 @@ export function LensEditConfigurationFlyout({
 
     return !visualizationStateIsEqual;
   }, [
-    attributes.references,
+    datasourceStates,
     datasourceId,
     datasourceMap,
+    attributes.references,
+    visualization.state,
     isNewPanel,
-    datasourceStates,
     visualizationMap,
     annotationGroups,
-    visualization.state,
   ]);
 
   const onCancel = useCallback(() => {

x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/state_helpers.ts

@@ -267,21 +267,24 @@ export async function initializeSources(
     options
   );
 
+  const initializedDatasourceStates = initializeDatasources({
+    datasourceMap,
+    datasourceStates,
+    initialContext,
+    indexPatternRefs,
+    indexPatterns,
+    references,
+  });
+
   return {
     indexPatterns,
     indexPatternRefs,
     annotationGroups,
-    datasourceStates: initializeDatasources({
-      datasourceMap,
-      datasourceStates,
-      initialContext,
-      indexPatternRefs,
-      indexPatterns,
-      references,
-    }),
+    datasourceStates: initializedDatasourceStates,
     visualizationState: initializeVisualization({
       visualizationMap,
       visualizationState,
+      datasourceStates,
       references,
       initialContext,
       annotationGroups,
@@ -292,11 +295,13 @@ export async function initializeSources(
 export function initializeVisualization({
   visualizationMap,
   visualizationState,
+  datasourceStates,
   references,
   annotationGroups,
 }: {
   visualizationState: VisualizationState;
   visualizationMap: VisualizationMap;
+  datasourceStates: DatasourceStates;
   references?: SavedObjectReference[];
   initialContext?: VisualizeFieldContext | VisualizeEditorContext;
   annotationGroups: Record<string, EventAnnotationGroupConfig>;
@@ -306,8 +311,9 @@ export function initializeVisualization({
       visualizationMap[visualizationState.activeId]?.initialize(
         () => '',
         visualizationState.state,
-        // initialize a new visualization with the color mapping off
         COLORING_METHOD,
+        datasourceStates,
+        // initialize a new visualization with the color mapping off
         annotationGroups,
         references
       ) ?? visualizationState.state
@@ -398,15 +404,6 @@ export async function persistedStateToExpression(
   );
 
   const visualization = visualizations[visualizationType!];
-  const activeVisualizationState = initializeVisualization({
-    visualizationMap: visualizations,
-    visualizationState: {
-      state: persistedVisualizationState,
-      activeId: visualizationType,
-    },
-    annotationGroups,
-    references: [...references, ...(internalReferences || [])],
-  });
   const datasourceStatesFromSO = Object.fromEntries(
     Object.entries(persistedDatasourceStates).map(([id, state]) => [
       id,
@@ -434,6 +431,17 @@ export async function persistedStateToExpression(
     indexPatternRefs,
   });
 
+  const activeVisualizationState = initializeVisualization({
+    visualizationMap: visualizations,
+    visualizationState: {
+      state: persistedVisualizationState,
+      activeId: visualizationType,
+    },
+    datasourceStates,
+    annotationGroups,
+    references: [...references, ...(internalReferences || [])],
+  });
+
   const datasourceLayers = getDatasourceLayers(datasourceStates, datasourceMap, indexPatterns);
 
   const datasourceId = getActiveDatasourceIdFromDoc(doc);

x-pack/platform/plugins/shared/lens/public/lens_ui_telemetry/color_telemetry_helpers.test.ts

@@ -18,8 +18,8 @@ const exampleAssignment = (
   valuesCount = 1,
   type = 'categorical',
   overrides = {}
-): ColorMapping.Config['assignments'][number] => {
-  const color: ColorMapping.Config['assignments'][number]['color'] =
+): ColorMapping.Assignment => {
+  const color: ColorMapping.Assignment['color'] =
     type === 'categorical'
       ? {
           type: 'categorical',
@@ -32,10 +32,10 @@ const exampleAssignment = (
         };
 
   return {
-    rule: {
-      type: 'matchExactly',
-      values: Array.from({ length: valuesCount }, () => faker.string.alpha()),
-    },
+    rules: Array.from({ length: valuesCount }, () => faker.string.alpha()).map((value) => ({
+      type: 'raw',
+      value,
+    })),
     color,
     touched: false,
     ...overrides,
@@ -51,9 +51,11 @@ const MANUAL_COLOR_MAPPING_CONFIG: ColorMapping.Config = {
   ],
   specialAssignments: [
     {
-      rule: {
-        type: 'other',
-      },
+      rules: [
+        {
+          type: 'other',
+        },
+      ],
       color: {
         type: 'categorical',
         paletteId: KbnPalette.ElasticClassic,

x-pack/platform/plugins/shared/lens/public/lens_ui_telemetry/color_telemetry_helpers.ts

@@ -92,10 +92,7 @@ const getUnassignedTermsType = (
 };
 
 const getTotalTermsCount = (assignments: ColorMapping.Config['assignments']) =>
-  assignments.reduce(
-    (acc, cur) => ('values' in cur.rule ? acc + cur.rule.values.length : acc + 1),
-    0
-  );
+  assignments.reduce((acc, { rules }) => acc + rules.length, 0);
 
 const getAvgCountTermsPerColor = (
   assignments: ColorMapping.Config['assignments'],

x-pack/platform/plugins/shared/lens/public/react_embeddable/helper.test.ts

@@ -6,8 +6,13 @@
  */
 import { BehaviorSubject } from 'rxjs';
 import { defaultDoc } from '../mocks/services_mock';
-import { deserializeState } from './helper';
+import { deserializeState, getStructuredDatasourceStates } from './helper';
 import { makeEmbeddableServices } from './mocks';
+import { FormBasedPersistedState } from '../datasources/form_based/types';
+import { TextBasedPersistedState } from '../datasources/form_based/esql_layer/types';
+import expect from 'expect';
+import { DatasourceState } from '../state_management';
+import { StructuredDatasourceStates } from './types';
 
 describe('Embeddable helpers', () => {
   describe('deserializeState', () => {
@@ -114,4 +119,51 @@ describe('Embeddable helpers', () => {
       });
     });
   });
+
+  describe('getStructuredDatasourceStates', () => {
+    const formBasedDSStateMock: FormBasedPersistedState = {
+      layers: {},
+    };
+    const textBasedDSStateMock: TextBasedPersistedState = {
+      layers: {},
+    };
+
+    it('should return structured datasourceStates from unknown datasourceStates', () => {
+      const mockDatasourceStates: Record<string, unknown> = {
+        formBased: formBasedDSStateMock,
+        textBased: textBasedDSStateMock,
+        other: textBasedDSStateMock,
+      };
+      const result = getStructuredDatasourceStates(mockDatasourceStates);
+
+      expect(result.formBased).toEqual(formBasedDSStateMock);
+      expect(result.textBased).toEqual(textBasedDSStateMock);
+      expect('other' in result).toBe(false);
+    });
+
+    it('should return structured datasourceStates from nested unknown datasourceStates', () => {
+      const wrap = (ds: unknown) => ({ state: ds, isLoading: false } satisfies DatasourceState);
+      const mockDatasourceStates: Record<string, unknown> = {
+        formBased: wrap(formBasedDSStateMock),
+        textBased: wrap(textBasedDSStateMock),
+        other: wrap(textBasedDSStateMock),
+      };
+      const result = getStructuredDatasourceStates(mockDatasourceStates);
+
+      expect(result.formBased).toEqual(formBasedDSStateMock);
+      expect(result.textBased).toEqual(textBasedDSStateMock);
+      expect('other' in result).toBe(false);
+    });
+
+    it('should return structured datasourceStates from structured datasourceStates', () => {
+      const mockDatasourceStates: StructuredDatasourceStates = {
+        formBased: formBasedDSStateMock,
+        textBased: textBasedDSStateMock,
+      };
+      const result = getStructuredDatasourceStates(mockDatasourceStates);
+
+      expect(result.formBased).toEqual(formBasedDSStateMock);
+      expect(result.textBased).toEqual(textBasedDSStateMock);
+    });
+  });
 });

x-pack/platform/plugins/shared/lens/public/react_embeddable/helper.ts

@@ -19,8 +19,16 @@ import fastIsEqual from 'fast-deep-equal';
 import { isOfAggregateQueryType } from '@kbn/es-query';
 import { RenderMode } from '@kbn/expressions-plugin/common';
 import { SavedObjectReference } from '@kbn/core/types';
-import type { LensEmbeddableStartServices, LensRuntimeState, LensSerializedState } from './types';
+import type {
+  LensEmbeddableStartServices,
+  LensRuntimeState,
+  LensSerializedState,
+  StructuredDatasourceStates,
+} from './types';
 import { loadESQLAttributes } from './esql';
+import { DatasourceStates, GeneralDatasourceStates } from '../state_management';
+import { FormBasedPersistedState } from '../datasources/form_based/types';
+import { TextBasedPersistedState } from '../datasources/form_based/esql_layer/types';
 
 export function createEmptyLensState(
   visualizationType: null | string = null,
@@ -172,3 +180,16 @@ export function extractInheritedViewModeObservable(
   }
   return new BehaviorSubject<ViewMode>('view');
 }
+
+export function getStructuredDatasourceStates(
+  datasourceStates?: Readonly<GeneralDatasourceStates>
+): StructuredDatasourceStates {
+  return {
+    formBased: ((datasourceStates as DatasourceStates)?.formBased?.state ??
+      datasourceStates?.formBased ??
+      undefined) as FormBasedPersistedState,
+    textBased: ((datasourceStates as DatasourceStates)?.textBased?.state ??
+      datasourceStates?.textBased ??
+      undefined) as TextBasedPersistedState,
+  };
+}

x-pack/platform/plugins/shared/lens/public/react_embeddable/types.ts

@@ -474,18 +474,20 @@ export interface ExpressionWrapperProps {
 
 export type GetStateType = () => LensRuntimeState;
 
-/**
- * Custom Lens component exported by the plugin
- * For better DX of Lens component consumers, expose a typed version of the serialized state
- */
+export interface StructuredDatasourceStates {
+  formBased?: FormBasedPersistedState;
+  textBased?: TextBasedPersistedState;
+}
 
-/** Utility function to build typed version for each chart */
+/** Utility type to build typed version for each chart */
 type TypedLensAttributes<TVisType, TVisState> = Simplify<
   Omit<LensDocument, 'savedObjectId' | 'type' | 'state' | 'visualizationType'> & {
     visualizationType: TVisType;
     state: Simplify<
       Omit<LensDocument['state'], 'datasourceStates' | 'visualization'> & {
         datasourceStates: {
+          // This is of type StructuredDatasourceStates but does not conform to Record<string, unknown>
+          // so I am leaving this alone until we improve this datasource typing structure.
           formBased?: FormBasedPersistedState;
           textBased?: TextBasedPersistedState;
         };

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/converter.test.ts

@@ -0,0 +1,427 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { convertToRawColorMappings, isDeprecatedColorMapping } from './converter';
+import { DeprecatedColorMappingConfig } from './types';
+import { DataType } from '../../../types';
+import { ColorMapping } from '@kbn/coloring';
+import { GenericIndexPatternColumn } from '../../../async_services';
+import { SerializedRangeKey } from '@kbn/data-plugin/common/search';
+
+type OldAssignment = DeprecatedColorMappingConfig['assignments'][number];
+type OldRule = OldAssignment['rule'];
+type OldSpecialRule = DeprecatedColorMappingConfig['specialAssignments'][number]['rule'];
+
+const baseConfig = {
+  assignments: [],
+  specialAssignments: [],
+  paletteId: 'default',
+  colorMode: {
+    type: 'categorical',
+  },
+} satisfies DeprecatedColorMappingConfig | ColorMapping.Config;
+
+const getOldColorMapping = (
+  rules: OldRule[],
+  specialRules: OldSpecialRule[] = []
+): DeprecatedColorMappingConfig => ({
+  assignments: rules.map((rule, i) => ({
+    rule,
+    color: {
+      type: 'categorical',
+      paletteId: 'default',
+      colorIndex: i,
+    },
+    touched: false,
+  })),
+  specialAssignments: specialRules.map((rule) => ({
+    rule,
+    color: {
+      type: 'loop',
+    },
+    touched: false,
+  })),
+  paletteId: 'default',
+  colorMode: {
+    type: 'categorical',
+  },
+});
+
+function buildOldColorMapping(
+  rules: OldRule[],
+  specialRules: OldSpecialRule[] = []
+): DeprecatedColorMappingConfig {
+  return getOldColorMapping(rules, specialRules);
+}
+
+describe('converter', () => {
+  describe('#convertToRawColorMappings', () => {
+    it('should convert config with no assignments', () => {
+      const oldConfig = buildOldColorMapping([]);
+      const newConfig = convertToRawColorMappings(oldConfig);
+      expect(newConfig.assignments).toHaveLength(0);
+    });
+
+    it('should keep top-level config', () => {
+      const oldConfig = buildOldColorMapping([]);
+      const newConfig = convertToRawColorMappings(oldConfig);
+      expect(newConfig).toMatchObject({
+        paletteId: 'default',
+        colorMode: {
+          type: 'categorical',
+        },
+      });
+    });
+
+    describe('type - auto', () => {
+      it('should convert single auto rule', () => {
+        const oldConfig = buildOldColorMapping([{ type: 'auto' }]);
+        const newConfig = convertToRawColorMappings(oldConfig);
+        expect(newConfig.assignments).toHaveLength(1);
+        expect(newConfig.assignments[0].color).toBeDefined();
+        expect(newConfig.assignments[0].rules).toEqual([]);
+      });
+
+      it('should convert multiple auto rule', () => {
+        const oldConfig = buildOldColorMapping([
+          { type: 'auto' },
+          { type: 'matchExactly', values: [] },
+          { type: 'auto' },
+        ]);
+        const newConfig = convertToRawColorMappings(oldConfig);
+        expect(newConfig.assignments).toHaveLength(3);
+        expect(newConfig.assignments[0].rules).toEqual([]);
+        expect(newConfig.assignments[2].rules).toEqual([]);
+      });
+    });
+
+    describe('type - matchExactly', () => {
+      type ExpectedRule = Partial<ColorMapping.RuleMatchRaw | ColorMapping.RuleMatch>;
+      interface ExpectedRulesByType {
+        types: Array<DataType | 'undefined'>;
+        expectedRule: ExpectedRule;
+      }
+      type MatchExactlyTestCase = [
+        oldStringValue: string | string[],
+        defaultExpectedRule: ExpectedRule,
+        expectedRulesByType: ExpectedRulesByType[]
+      ];
+
+      const buildOldColorMappingFromValues = (values: Array<string | string[]>) =>
+        buildOldColorMapping([{ type: 'matchExactly', values }]);
+
+      it('should handle missing column', () => {
+        const oldConfig = buildOldColorMapping([{ type: 'matchExactly', values: ['test'] }]);
+        const newConfig = convertToRawColorMappings(oldConfig, undefined);
+        expect(newConfig.assignments).toHaveLength(1);
+        expect(newConfig.assignments[0].rules).toHaveLength(1);
+        expect(newConfig.assignments[0].rules[0]).toEqual({
+          type: 'match',
+          pattern: 'test',
+          matchEntireWord: true,
+          matchCase: true,
+        });
+      });
+
+      describe('multi_terms', () => {
+        it('should convert array of string values as MultiFieldKey', () => {
+          const values: string[] = ['some-string', '123', '0', '1', '1744261200000', '__other__'];
+          const oldConfig = buildOldColorMappingFromValues([values]);
+          const newConfig = convertToRawColorMappings(oldConfig, {});
+          const rule = newConfig.assignments[0].rules[0];
+
+          expect(rule).toEqual({
+            type: 'raw',
+            value: {
+              keys: ['some-string', '123', '0', '1', '1744261200000', '__other__'],
+              type: 'multiFieldKey',
+            },
+          });
+        });
+
+        it('should convert array of strings in multi_terms as MultiFieldKey', () => {
+          const oldConfig = buildOldColorMappingFromValues([['some-string']]);
+          const newConfig = convertToRawColorMappings(oldConfig, { fieldType: 'multi_terms' });
+          const rule = newConfig.assignments[0].rules[0];
+
+          expect(rule).toEqual({
+            type: 'raw',
+            value: {
+              keys: ['some-string'],
+              type: 'multiFieldKey',
+            },
+          });
+        });
+
+        it('should convert single string as basic match even in multi_terms column', () => {
+          const oldConfig = buildOldColorMappingFromValues(['some-string']);
+          const newConfig = convertToRawColorMappings(oldConfig, { fieldType: 'multi_terms' });
+          const rule = newConfig.assignments[0].rules[0];
+
+          expect(rule).toEqual({
+            type: 'match',
+            pattern: 'some-string',
+            matchEntireWord: true,
+            matchCase: true,
+          });
+        });
+      });
+
+      describe('range', () => {
+        it.each<[rangeString: string, expectedRange: Pick<SerializedRangeKey, 'from' | 'to'>]>([
+          ['from:0,to:1000', { from: 0, to: 1000 }],
+          ['from:-1000,to:1000', { from: -1000, to: 1000 }],
+          ['from:-1000,to:0', { from: -1000, to: 0 }],
+          ['from:1000,to:undefined', { from: 1000, to: null }],
+          ['from:undefined,to:1000', { from: null, to: 1000 }],
+          ['from:undefined,to:undefined', { from: null, to: null }],
+        ])('should convert range string %j to RangeKey', (rangeString, expectedRange) => {
+          const oldConfig = buildOldColorMappingFromValues([rangeString]);
+          const newConfig = convertToRawColorMappings(oldConfig, { fieldType: 'range' });
+          const rule = newConfig.assignments[0].rules[0];
+
+          expect(rule).toEqual({
+            type: 'raw',
+            value: {
+              ...expectedRange,
+              type: 'rangeKey',
+              ranges: [],
+            },
+          });
+        });
+
+        it('should convert non-range string to match', () => {
+          const oldConfig = buildOldColorMappingFromValues(['not-a-range']);
+          const newConfig = convertToRawColorMappings(oldConfig, { fieldType: 'range' });
+          const rule = newConfig.assignments[0].rules[0];
+
+          expect(rule).toEqual({
+            type: 'match',
+            pattern: 'not-a-range',
+            matchEntireWord: true,
+            matchCase: true,
+          });
+        });
+      });
+
+      describe.each<DataType | undefined>([
+        'number',
+        'boolean',
+        'date',
+        'string',
+        'ip',
+        undefined, // other
+      ])('Column dataType - %s', (dataType) => {
+        const column: Partial<GenericIndexPatternColumn> = { dataType };
+
+        it.each<MatchExactlyTestCase>([
+          [
+            '123.456',
+            {
+              type: 'raw',
+              value: 123.456,
+            },
+            [
+              { types: ['string', 'ip'], expectedRule: { type: 'raw', value: '123.456' } },
+              {
+                types: ['undefined', 'boolean'],
+                expectedRule: { type: 'match', pattern: '123.456' },
+              },
+            ],
+          ],
+          [
+            '1744261200000',
+            {
+              type: 'raw',
+              value: 1744261200000,
+            },
+            [
+              { types: ['string', 'ip'], expectedRule: { type: 'raw', value: '1744261200000' } },
+              {
+                types: ['undefined', 'boolean'],
+                expectedRule: { type: 'match', pattern: '1744261200000' },
+              },
+            ],
+          ],
+          [
+            '__other__',
+            {
+              type: 'raw',
+              value: '__other__',
+            },
+            [{ types: ['undefined'], expectedRule: { type: 'match', pattern: '__other__' } }],
+          ],
+          [
+            'some-string',
+            {
+              type: 'raw',
+              value: 'some-string',
+            },
+            [
+              {
+                types: ['undefined', 'number', 'boolean', 'date'],
+                expectedRule: { type: 'match', pattern: 'some-string' },
+              },
+            ],
+          ],
+          [
+            'false',
+            {
+              type: 'raw',
+              value: 'false',
+            },
+            [
+              {
+                types: ['undefined', 'number', 'date'],
+                expectedRule: { type: 'match', pattern: 'false' },
+              },
+            ],
+          ],
+          [
+            'true',
+            {
+              type: 'raw',
+              value: 'true',
+            },
+            [
+              {
+                types: ['undefined', 'number', 'date'],
+                expectedRule: { type: 'match', pattern: 'true' },
+              },
+            ],
+          ],
+          [
+            '0', // false
+            {
+              type: 'raw',
+              value: 0,
+            },
+            [
+              { types: ['string', 'ip'], expectedRule: { type: 'raw', value: '0' } },
+              { types: ['undefined'], expectedRule: { type: 'match', pattern: '0' } },
+            ],
+          ],
+          [
+            '1', // true
+            {
+              type: 'raw',
+              value: 1,
+            },
+            [
+              { types: ['string', 'ip'], expectedRule: { type: 'raw', value: '1' } },
+              { types: ['undefined'], expectedRule: { type: 'match', pattern: '1' } },
+            ],
+          ],
+          [
+            '127.0.0.1',
+            {
+              type: 'raw',
+              value: '127.0.0.1',
+            },
+            [
+              {
+                types: ['undefined', 'number', 'boolean', 'date'],
+                expectedRule: { type: 'match', pattern: '127.0.0.1' },
+              },
+            ],
+          ],
+        ])('should correctly convert %j', (value, defaultExpectedRule, expectedRulesByType) => {
+          const oldConfig = buildOldColorMappingFromValues([value]);
+          const expectedRule =
+            expectedRulesByType.find((r) => r.types.includes(dataType ?? 'undefined'))
+              ?.expectedRule ?? defaultExpectedRule;
+          const newConfig = convertToRawColorMappings(oldConfig, column);
+          const rule = newConfig.assignments[0].rules[0];
+
+          if (expectedRule.type === 'match') {
+            // decorate match type with default constants
+            expectedRule.matchEntireWord = true;
+            expectedRule.matchCase = true;
+          }
+
+          expect(rule).toEqual(expectedRule);
+        });
+      });
+    });
+  });
+
+  describe('#isDeprecatedColorMapping', () => {
+    const baseAssignment = {
+      color: {
+        type: 'categorical',
+        paletteId: 'default',
+        colorIndex: 3,
+      },
+      touched: false,
+    } satisfies Omit<OldAssignment, 'rule'>;
+
+    it('should return true if assignments.rule exists', () => {
+      const isDeprecated = isDeprecatedColorMapping({
+        ...baseConfig,
+        assignments: [
+          {
+            ...baseAssignment,
+            rule: {
+              type: 'auto',
+            },
+          },
+        ],
+      });
+      expect(isDeprecated).toBe(true);
+    });
+
+    it('should return true if specialAssignments.rule exists', () => {
+      const isDeprecated = isDeprecatedColorMapping({
+        ...baseConfig,
+        specialAssignments: [
+          {
+            ...baseAssignment,
+            rule: {
+              type: 'other',
+            },
+          },
+        ],
+      });
+      expect(isDeprecated).toBe(true);
+    });
+
+    it('should return false if assignments.rule does not exist', () => {
+      const isDeprecated = isDeprecatedColorMapping({
+        ...baseConfig,
+        assignments: [
+          {
+            ...baseAssignment,
+            rules: [
+              {
+                type: 'match',
+                pattern: 'test',
+              },
+            ],
+          },
+        ],
+      });
+      expect(isDeprecated).toBe(false);
+    });
+
+    it('should return false if specialAssignments.rule does not exist', () => {
+      const isDeprecated = isDeprecatedColorMapping({
+        ...baseConfig,
+        specialAssignments: [
+          {
+            ...baseAssignment,
+            rules: [
+              {
+                type: 'other',
+              },
+            ],
+          },
+        ],
+      });
+      expect(isDeprecated).toBe(false);
+    });
+  });
+});

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/converter.ts

@@ -0,0 +1,160 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ColorMapping } from '@kbn/coloring';
+import { MultiFieldKey, RangeKey, SerializedValue } from '@kbn/data-plugin/common';
+import { DeprecatedColorMappingConfig } from './types';
+import { ColumnMeta } from './utils';
+
+/**
+ * Converts old stringified colorMapping configs to new raw value configs
+ */
+export function convertToRawColorMappings(
+  colorMapping: DeprecatedColorMappingConfig | ColorMapping.Config,
+  columnMeta?: ColumnMeta | null
+): ColorMapping.Config {
+  return {
+    ...colorMapping,
+    assignments: colorMapping.assignments.map((oldAssignment) => {
+      if (isValidColorMappingAssignment(oldAssignment)) return oldAssignment;
+      return convertColorMappingAssignment(oldAssignment, columnMeta);
+    }),
+    specialAssignments: colorMapping.specialAssignments.map((oldAssignment) => {
+      if (isValidColorMappingAssignment(oldAssignment)) return oldAssignment;
+      return {
+        color: oldAssignment.color,
+        touched: oldAssignment.touched,
+        rules: [oldAssignment.rule],
+      };
+    }),
+  };
+}
+
+function convertColorMappingAssignment(
+  oldAssignment: DeprecatedColorMappingConfig['assignments'][number],
+  columnMeta?: ColumnMeta | null
+): ColorMapping.Assignment {
+  return {
+    color: oldAssignment.color,
+    touched: oldAssignment.touched,
+    rules: convertColorMappingRule(oldAssignment.rule, columnMeta),
+  };
+}
+
+const NO_VALUE = Symbol('no-value');
+
+function convertColorMappingRule(
+  rule: DeprecatedColorMappingConfig['assignments'][number]['rule'],
+  columnMeta?: ColumnMeta | null
+): ColorMapping.ColorRule[] {
+  switch (rule.type) {
+    case 'auto':
+      return [];
+    case 'matchExactly':
+      return rule.values.map((value) => {
+        const rawValue = convertToRawValue(value, columnMeta);
+
+        if (rawValue !== NO_VALUE) {
+          return {
+            type: 'raw',
+            value: rawValue,
+          };
+        }
+
+        return {
+          type: 'match',
+          pattern: String(value),
+          matchEntireWord: true,
+          matchCase: true,
+        };
+      });
+
+    // Rules below not yet used, adding conversions for completeness
+    case 'matchExactlyCI':
+      return rule.values.map((value) => ({
+        type: 'match',
+        pattern: Array.isArray(value) ? value.join(' ') : value,
+        matchEntireWord: true,
+        matchCase: false,
+      }));
+    case 'regex':
+      return [{ type: rule.type, pattern: rule.values }];
+    case 'range':
+    default:
+      return [rule];
+  }
+}
+
+/**
+ * Attempts to convert the previously stringified raw values into their raw/serialized form
+ *
+ * Note: we use the `NO_VALUE` symbol to avoid collisions with falsy raw values
+ */
+function convertToRawValue(
+  value: string | string[],
+  columnMeta?: ColumnMeta | null
+): SerializedValue | symbol {
+  if (!columnMeta) return NO_VALUE;
+
+  // all array values are multi-term
+  if (columnMeta.fieldType === 'multi_terms' || Array.isArray(value)) {
+    if (typeof value === 'string') return NO_VALUE; // cannot assume this as multi-field
+    return new MultiFieldKey({ key: value }).serialize();
+  }
+
+  if (columnMeta.fieldType === 'range') {
+    return RangeKey.isRangeKeyString(value) ? RangeKey.fromString(value).serialize() : NO_VALUE;
+  }
+
+  switch (columnMeta.dataType) {
+    case 'boolean':
+      if (value === '__other__' || value === 'true' || value === 'false') return value; // bool could have __other__ as a string
+      if (value === '0' || value === '1') return Number(value);
+      break;
+    case 'number':
+    case 'date':
+      if (value === '__other__') return value; // numbers can have __other__ as a string
+      const numberValue = Number(value);
+      if (isFinite(numberValue)) return numberValue;
+      break;
+    case 'string':
+    case 'ip':
+      return value; // unable to distinguish manually added values
+    default:
+      return NO_VALUE; // treat all other other dataType as custom match string values
+  }
+  return NO_VALUE;
+}
+
+function isValidColorMappingAssignment<
+  T extends
+    | DeprecatedColorMappingConfig['assignments'][number]
+    | DeprecatedColorMappingConfig['specialAssignments'][number]
+    | ColorMapping.Config['assignments'][number]
+    | ColorMapping.Config['specialAssignments'][number]
+>(
+  assignment: T
+): assignment is Exclude<
+  T,
+  | DeprecatedColorMappingConfig['assignments'][number]
+  | DeprecatedColorMappingConfig['specialAssignments'][number]
+> {
+  return 'rules' in assignment;
+}
+
+export function isDeprecatedColorMapping<
+  T extends DeprecatedColorMappingConfig | ColorMapping.Config
+>(colorMapping?: T): colorMapping is Exclude<T, ColorMapping.Config> {
+  if (!colorMapping) return false;
+  return Boolean(
+    colorMapping.assignments &&
+      (colorMapping.assignments.some((assignment) => !isValidColorMappingAssignment(assignment)) ||
+        colorMapping.specialAssignments.some(
+          (specialAssignment) => !isValidColorMappingAssignment(specialAssignment)
+        ))
+  );
+}

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/index.ts

@@ -0,0 +1,10 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export type { DeprecatedColorMappingConfig } from './types';
+export { convertToRawColorMappings, isDeprecatedColorMapping } from './converter';
+export { getColumnMetaFn } from './utils';

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/types.ts

@@ -0,0 +1,117 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/** @deprecated */
+interface DeprecatedColorMappingColorCode {
+  type: 'colorCode';
+  colorCode: string;
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingCategoricalColor {
+  type: 'categorical';
+  paletteId: string;
+  colorIndex: number;
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingGradientColor {
+  type: 'gradient';
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingLoopColor {
+  type: 'loop';
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingRuleAuto {
+  type: 'auto';
+}
+/** @deprecated */
+interface DeprecatedColorMappingRuleMatchExactly {
+  type: 'matchExactly';
+  values: Array<string | string[]>;
+}
+/** @deprecated */
+interface DeprecatedColorMappingRuleMatchExactlyCI {
+  type: 'matchExactlyCI';
+  values: string[];
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingRuleRange {
+  type: 'range';
+  min: number;
+  max: number;
+  minInclusive: boolean;
+  maxInclusive: boolean;
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingRuleRegExp {
+  type: 'regex';
+  values: string;
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingRuleOthers {
+  type: 'other';
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingAssignment<R, C> {
+  rule: R;
+  color: C;
+  touched: boolean;
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingCategoricalColorMode {
+  type: 'categorical';
+}
+
+/** @deprecated */
+interface DeprecatedColorMappingGradientColorMode {
+  type: 'gradient';
+  steps: Array<
+    (DeprecatedColorMappingCategoricalColor | DeprecatedColorMappingColorCode) & {
+      touched: boolean;
+    }
+  >;
+  sort: 'asc' | 'desc';
+}
+
+/**
+ * Old color mapping state meant for type safety during runtime migrations of old configurations
+ *
+ * @deprecated Use `ColorMapping.Config`
+ */
+export interface DeprecatedColorMappingConfig {
+  paletteId: string;
+  colorMode: DeprecatedColorMappingCategoricalColorMode | DeprecatedColorMappingGradientColorMode;
+  assignments: Array<
+    DeprecatedColorMappingAssignment<
+      | DeprecatedColorMappingRuleAuto
+      | DeprecatedColorMappingRuleMatchExactly
+      | DeprecatedColorMappingRuleMatchExactlyCI
+      | DeprecatedColorMappingRuleRange
+      | DeprecatedColorMappingRuleRegExp,
+      | DeprecatedColorMappingCategoricalColor
+      | DeprecatedColorMappingColorCode
+      | DeprecatedColorMappingGradientColor
+    >
+  >;
+  specialAssignments: Array<
+    DeprecatedColorMappingAssignment<
+      DeprecatedColorMappingRuleOthers,
+      | DeprecatedColorMappingCategoricalColor
+      | DeprecatedColorMappingColorCode
+      | DeprecatedColorMappingLoopColor
+    >
+  >;
+}

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/utils.test.ts

@@ -0,0 +1,152 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+  TextBasedLayer,
+  TextBasedPersistedState,
+} from '../../../datasources/form_based/esql_layer/types';
+import { FormBasedLayer, FormBasedPersistedState } from '../../../datasources/form_based/types';
+import { StructuredDatasourceStates } from '../../../react_embeddable/types';
+import { ColumnMeta, getColumnMetaFn } from './utils';
+
+const layerId = 'layer-1';
+const columnId = 'column-1';
+
+const getDatasourceStatesMock = (
+  type: keyof StructuredDatasourceStates,
+  dataType?: ColumnMeta['dataType'],
+  fieldType?: ColumnMeta['fieldType']
+) => {
+  if (type === 'formBased') {
+    return {
+      formBased: {
+        layers: {
+          [layerId]: {
+            columns: {
+              [columnId]: {
+                dataType,
+                params: {
+                  parentFormat: { id: fieldType },
+                },
+              },
+            },
+          } as unknown as FormBasedLayer,
+        },
+      } satisfies FormBasedPersistedState,
+    };
+  }
+
+  if (type === 'textBased') {
+    return {
+      textBased: {
+        layers: {
+          [layerId]: {
+            columns: [{ columnId, meta: { type: dataType } }],
+          } as unknown as TextBasedLayer,
+        },
+      } satisfies TextBasedPersistedState,
+    };
+  }
+};
+
+describe('utils', () => {
+  describe('getColumnMetaFn', () => {
+    const mockDataType = 'string';
+    const mockFieldType = 'terms';
+
+    it('should return null if neither type exists', () => {
+      const mockDatasourceState = { 'not-supported': {} };
+      const resultFn = getColumnMetaFn(mockDatasourceState);
+
+      expect(resultFn).toBeNull();
+    });
+
+    describe('formBased datasourceState', () => {
+      it('should correct dataType and fieldType', () => {
+        const mockDatasourceState = getDatasourceStatesMock(
+          'formBased',
+          mockDataType,
+          mockFieldType
+        );
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, columnId);
+
+        expect(result.dataType).toBe(mockDataType);
+        expect(result.fieldType).toBe(mockFieldType);
+      });
+
+      it('should undefined dataType and fieldType if column not found', () => {
+        const mockDatasourceState = getDatasourceStatesMock('formBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, 'bad-column');
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+
+      it('should undefined dataType and fieldType if layer not found', () => {
+        const mockDatasourceState = getDatasourceStatesMock('formBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn('bad-layer', columnId);
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+
+      it('should undefined dataType and fieldType if missing', () => {
+        const mockDatasourceState = getDatasourceStatesMock('formBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, columnId);
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+    });
+
+    describe('textBased datasourceState', () => {
+      it('should correct dataType and fieldType', () => {
+        const mockDatasourceState = getDatasourceStatesMock(
+          'textBased',
+          mockDataType,
+          mockFieldType
+        );
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, columnId);
+
+        expect(result.dataType).toBe(mockDataType);
+        expect(result.fieldType).toBeUndefined(); // no fieldType needed for textBased
+      });
+
+      it('should undefined dataType and fieldType if column not found', () => {
+        const mockDatasourceState = getDatasourceStatesMock('textBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, 'bad-column');
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+
+      it('should undefined dataType and fieldType if layer not found', () => {
+        const mockDatasourceState = getDatasourceStatesMock('textBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn('bad-layer', columnId);
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+
+      it('should undefined dataType and fieldType if missing', () => {
+        const mockDatasourceState = getDatasourceStatesMock('textBased');
+        const resultFn = getColumnMetaFn(mockDatasourceState)!;
+        const result = resultFn(layerId, columnId);
+
+        expect(result.dataType).toBeUndefined();
+        expect(result.fieldType).toBeUndefined();
+      });
+    });
+  });
+});

x-pack/platform/plugins/shared/lens/public/runtime_state/converters/raw_color_mappings/utils.ts

@@ -0,0 +1,51 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { DatatableColumnType } from '@kbn/expressions-plugin/common';
+import { GenericIndexPatternColumn } from '../../../datasources/form_based/types';
+import { getStructuredDatasourceStates } from '../../../react_embeddable/helper';
+import { GeneralDatasourceStates } from '../../../state_management';
+
+export interface ColumnMeta {
+  fieldType?: string | 'multi_terms' | 'range';
+  dataType?: GenericIndexPatternColumn['dataType'] | DatatableColumnType;
+}
+
+export function getColumnMetaFn(
+  datasourceStates?: Readonly<GeneralDatasourceStates>
+): ((layerId: string, columnId: string) => ColumnMeta) | null {
+  const datasources = getStructuredDatasourceStates(datasourceStates);
+
+  if (datasources.formBased?.layers) {
+    const layers = datasources.formBased.layers;
+
+    return (layerId, columnId) => {
+      const column = layers[layerId]?.columns?.[columnId];
+      return {
+        fieldType:
+          column && 'params' in column
+            ? (column.params as { parentFormat?: { id?: string } })?.parentFormat?.id
+            : undefined,
+        dataType: column?.dataType,
+      };
+    };
+  }
+
+  if (datasources.textBased?.layers) {
+    const layers = datasources.textBased.layers;
+
+    return (layerId, columnId) => {
+      const column = layers[layerId]?.columns?.find((c) => c.columnId === columnId);
+
+      return {
+        dataType: column?.meta?.type,
+      };
+    };
+  }
+
+  return null;
+}

x-pack/platform/plugins/shared/lens/public/runtime_state/index.ts

@@ -0,0 +1,8 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export * from './converters/raw_color_mappings';

x-pack/platform/plugins/shared/lens/public/shared_components/coloring/color_mapping_by_terms.tsx

@@ -12,9 +12,11 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
   EuiFormRow,
+  EuiIconTip,
   EuiSpacer,
   EuiSwitch,
   EuiText,
+  useEuiTheme,
 } from '@elastic/eui';
 import {
   ColorMapping,
@@ -27,6 +29,8 @@ import {
 } from '@kbn/coloring';
 import { i18n } from '@kbn/i18n';
 import { KbnPalettes } from '@kbn/palettes';
+import { IFieldFormat } from '@kbn/field-formats-plugin/common';
+import { SerializedValue } from '@kbn/data-plugin/common';
 import { trackUiCounterEvents } from '../../lens_ui_telemetry';
 import { PalettePicker } from '../palette_picker';
 import { PalettePanelContainer } from './palette_panel_container';
@@ -42,7 +46,9 @@ interface ColorMappingByTermsProps {
   setColorMapping: (colorMapping?: ColorMapping.Config) => void;
   paletteService: PaletteRegistry;
   panelRef: MutableRefObject<HTMLDivElement | null>;
-  categories: Array<string | string[]>;
+  categories: SerializedValue[];
+  formatter?: IFieldFormat;
+  allowCustomMatch?: boolean;
 }
 
 export function ColorMappingByTerms({
@@ -56,7 +62,10 @@ export function ColorMappingByTerms({
   paletteService,
   panelRef,
   categories,
+  formatter,
+  allowCustomMatch,
 }: ColorMappingByTermsProps) {
+  const { euiTheme } = useEuiTheme();
   const [useNewColorMapping, setUseNewColorMapping] = useState(Boolean(colorMapping));
 
   return (
@@ -101,6 +110,20 @@ export function ColorMappingByTerms({
                       {i18n.translate('xpack.lens.colorMapping.tryLabel', {
                         defaultMessage: 'Use the new Color Mapping feature',
                       })}{' '}
+                      {(colorMapping?.assignments.length ?? 0) > 0 && (
+                        <EuiIconTip
+                          content={i18n.translate(
+                            'xpack.lens.colorMapping.helpIncompatibleFieldDotLabel',
+                            {
+                              defaultMessage: 'Disabling Color Mapping will clear all assignments',
+                            }
+                          )}
+                          position="top"
+                          size="s"
+                          type="dot"
+                          color={euiTheme.colors.warning}
+                        />
+                      )}{' '}
                       <EuiBadge color="hollow">
                         {i18n.translate('xpack.lens.colorMapping.techPreviewLabel', {
                           defaultMessage: 'Tech preview',
@@ -128,6 +151,8 @@ export function ColorMappingByTerms({
                   onModelUpdate={setColorMapping}
                   specialTokens={SPECIAL_TOKENS_STRING_CONVERSION}
                   palettes={palettes}
+                  formatter={formatter}
+                  allowCustomMatch={allowCustomMatch}
                   data={{
                     type: 'categories',
                     categories,

x-pack/platform/plugins/shared/lens/public/shared_components/coloring/get_cell_color_fn.ts

@@ -10,12 +10,14 @@ import {
   PaletteOutput,
   PaletteRegistry,
   getSpecialString,
+  getValueKey,
 } from '@kbn/coloring';
 import { CustomPaletteState } from '@kbn/charts-plugin/common';
 import { KbnPalettes } from '@kbn/palettes';
+import { RawValue } from '@kbn/data-plugin/common';
 import { getColorAccessorFn } from './color_mapping_accessor';
 
-export type CellColorFn = (value?: number | string | null) => string | null;
+export type CellColorFn = (value: RawValue) => string | null;
 
 export function getCellColorFn(
   paletteService: PaletteRegistry,
@@ -41,17 +43,17 @@ export function getCellColorFn(
     if (colorMapping) {
       return getColorAccessorFn(palettes, colorMapping, data, isDarkMode);
     } else if (palette) {
-      return (category) => {
-        if (category === undefined || category === null) return null;
+      return (value: RawValue) => {
+        if (value === undefined || value === null) return null;
 
-        const strCategory = String(category); // can be a number as a string
+        const key = getValueKey(value);
 
         return paletteService.get(palette.name).getCategoricalColor(
           [
             {
-              name: getSpecialString(strCategory), // needed to sync special categories (i.e. '')
+              name: getSpecialString(key), // needed to sync special categories (i.e. '')
               rankAtDepth: Math.max(
-                data.categories.findIndex((v) => v === strCategory),
+                data.categories.findIndex((v) => v === key),
                 0
               ),
               totalSeriesAtDepth: data.categories.length || 1,

x-pack/platform/plugins/shared/lens/public/state_management/__snapshots__/load_initial.test.tsx.snap

@@ -84,7 +84,7 @@ Object {
           "language": "lucene",
           "query": "",
         },
-        "visualization": Object {},
+        "visualization": "testVis initial state",
       },
       "title": "",
       "type": "lens",

x-pack/platform/plugins/shared/lens/public/state_management/init_middleware/load_initial.ts

@@ -286,7 +286,13 @@ async function loadFromSavedObject(
           : !inlineEditing
           ? data.search.session.start()
           : undefined,
-      persistedDoc: doc,
+      persistedDoc: {
+        ...doc,
+        state: {
+          ...doc.state,
+          visualization: visualizationState,
+        },
+      },
       activeDatasourceId: getInitialDatasourceId(loaderSharedArgs.datasourceMap, doc),
       visualization: {
         activeId: doc.visualizationType,

x-pack/platform/plugins/shared/lens/public/state_management/load_initial.test.tsx

@@ -234,7 +234,14 @@ describe('Initializing the store', () => {
 
       expect(store.getState()).toEqual({
         lens: expect.objectContaining({
-          persistedDoc: { ...defaultDoc, type: DOC_TYPE },
+          persistedDoc: expect.objectContaining({
+            ...defaultDoc,
+            type: DOC_TYPE,
+            state: {
+              ...defaultDoc.state,
+              visualization: 'testVis initial state',
+            },
+          }),
           query: defaultDoc.state.query,
           isLoading: false,
           activeDatasourceId: 'testDatasource',

x-pack/platform/plugins/shared/lens/public/state_management/types.ts

@@ -9,6 +9,7 @@ import type { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
 import type { EmbeddableEditorState } from '@kbn/embeddable-plugin/public';
 import type { AggregateQuery, Filter, Query } from '@kbn/es-query';
 import type { SavedQuery } from '@kbn/data-plugin/public';
+import { $Values } from 'utility-types';
 import type { MainHistoryLocationState } from '../../common/locator/locator';
 import type { LensDocument } from '../persistence';
 
@@ -24,6 +25,7 @@ import type {
   IndexPatternRef,
   AnnotationGroups,
 } from '../types';
+import { StructuredDatasourceStates } from '../react_embeddable/types';
 export interface VisualizationState {
   activeId: string | null;
   state: unknown;
@@ -34,12 +36,29 @@ export interface DataViewsState {
   indexPatterns: Record<string, IndexPattern>;
 }
 
-export interface DatasourceState {
+export interface DatasourceState<S = unknown> {
   isLoading: boolean;
-  state: unknown;
+  state: S;
 }
-
 export type DatasourceStates = Record<string, DatasourceState>;
+
+/**
+ * A type to encompass all variants of DatasourceState types
+ *
+ * TODO: cleanup types/structure of datasources
+ */
+export type GeneralDatasourceState =
+  | unknown
+  | $Values<StructuredDatasourceStates>
+  | DatasourceState<unknown>
+  | DatasourceState<$Values<StructuredDatasourceStates>>;
+/**
+ * A type to encompass all variants of DatasourceStates types
+ */
+export type GeneralDatasourceStates =
+  | Record<string, GeneralDatasourceState>
+  | StructuredDatasourceStates;
+
 export interface PreviewState {
   visualization: VisualizationState;
   datasourceStates: DatasourceStates;

x-pack/platform/plugins/shared/lens/public/state_management/utils.ts

@@ -13,7 +13,7 @@ export const getDatasourceLayers = memoizeOne(function getDatasourceLayers(
   datasourceStates: DatasourceStates,
   datasourceMap: DatasourceMap,
   indexPatterns: DataViewsState['indexPatterns']
-) {
+): DatasourceLayers {
   const datasourceLayers: DatasourceLayers = {};
   Object.keys(datasourceMap)
     .filter((id) => datasourceStates[id] && !datasourceStates[id].isLoading)

x-pack/platform/plugins/shared/lens/public/types.ts

@@ -60,7 +60,7 @@ import {
   LENS_EDIT_PAGESIZE_ACTION,
 } from './visualizations/datatable/components/constants';
 import type { LensInspector } from './lens_inspector_service';
-import type { DataViewsState } from './state_management/types';
+import type { DataViewsState, GeneralDatasourceStates } from './state_management/types';
 import type { IndexPatternServiceAPI } from './data_views_service/service';
 import type { LensDocument } from './persistence/saved_object_store';
 import { TableInspectorAdapter } from './editor_frame_service/types';
@@ -1073,12 +1073,14 @@ export interface Visualization<T = unknown, P = T, ExtraAppendLayerArg = unknown
     (
       addNewLayer: () => string,
       nonPersistedState?: T,
-      mainPalette?: SuggestionRequest['mainPalette']
+      mainPalette?: SuggestionRequest['mainPalette'],
+      datasourceStates?: GeneralDatasourceStates
     ): T;
     (
       addNewLayer: () => string,
       persistedState: P,
       mainPalette?: SuggestionRequest['mainPalette'],
+      datasourceStates?: GeneralDatasourceStates,
       annotationGroups?: AnnotationGroups,
       references?: SavedObjectReference[]
     ): T;

x-pack/platform/plugins/shared/lens/public/utils.test.ts

@@ -7,7 +7,7 @@
 
 import { createDatatableUtilitiesMock } from '@kbn/data-plugin/common/mocks';
 import { Datatable } from '@kbn/expressions-plugin/public';
-import { getUniqueLabelGenerator, inferTimeField, isLensRange, renewIDs } from './utils';
+import { getUniqueLabelGenerator, inferTimeField, renewIDs } from './utils';
 
 const datatableUtilities = createDatatableUtilitiesMock();
 
@@ -187,24 +187,4 @@ describe('utils', () => {
       expect([' ', ' '].map(labelGenerator)).toEqual(['[Untitled]', '[Untitled] [1]']);
     });
   });
-
-  describe('isRange', () => {
-    it.each<[expected: boolean, input: unknown]>([
-      [true, { from: 0, to: 100, label: '' }],
-      [true, { from: 0, to: null, label: '' }],
-      [true, { from: null, to: 100, label: '' }],
-      [false, { from: 0, to: 100 }],
-      [false, { from: 0, to: null }],
-      [false, { from: null, to: 100 }],
-      [false, { from: 0 }],
-      [false, { to: 100 }],
-      [false, null],
-      [false, undefined],
-      [false, 123],
-      [false, 'string'],
-      [false, {}],
-    ])('should return %s for %j', (expected, input) => {
-      expect(isLensRange(input)).toBe(expected);
-    });
-  });
 });

x-pack/platform/plugins/shared/lens/public/utils.ts

@@ -39,7 +39,6 @@ import {
 import type { DatasourceStates, VisualizationState } from './state_management';
 import type { IndexPatternServiceAPI } from './data_views_service/service';
 import { COLOR_MAPPING_OFF_BY_DEFAULT } from '../common/constants';
-import type { RangeTypeLens } from './datasources/form_based/operations/definitions/ranges';
 
 export function getVisualizeGeoFieldMessage(fieldType: string) {
   return i18n.translate('xpack.lens.visualizeGeoFieldMessage', {
@@ -48,17 +47,6 @@ export function getVisualizeGeoFieldMessage(fieldType: string) {
   });
 }
 
-export function isLensRange(range: unknown = {}): range is RangeTypeLens {
-  if (!range || typeof range !== 'object') return false;
-  const { from, to, label } = range as RangeTypeLens;
-
-  return (
-    label !== undefined &&
-    (typeof from === 'number' || from === null) &&
-    (typeof to === 'number' || to === null)
-  );
-}
-
 export function getResolvedDateRange(timefilter: TimefilterContract) {
   const { from, to } = timefilter.getTime();
   return { fromDate: from, toDate: to };

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

@@ -10,24 +10,13 @@ import { EuiDataGridCellValueElementProps, EuiLink } from '@elastic/eui';
 import classNames from 'classnames';
 import { PaletteOutput } from '@kbn/coloring';
 import { CustomPaletteState } from '@kbn/charts-plugin/common';
+import { RawValue } from '@kbn/data-plugin/common';
 import type { FormatFactory } from '../../../../common/types';
 import type { DatatableColumnConfig } from '../../../../common/expressions';
 import type { DataContextType } from './types';
 import { getContrastColor } from '../../../shared_components/coloring/utils';
 import { CellColorFn } from '../../../shared_components/coloring/get_cell_color_fn';
 
-import { isLensRange } from '../../../utils';
-
-const getParsedValue = (v: unknown) => {
-  if (v == null || typeof v === 'number') {
-    return v;
-  }
-  if (isLensRange(v)) {
-    return v.toString();
-  }
-  return String(v);
-};
-
 export const createGridCell = (
   formatters: Record<string, ReturnType<FormatFactory>>,
   columnConfig: DatatableColumnConfig,
@@ -42,8 +31,8 @@ export const createGridCell = (
 ) => {
   return ({ rowIndex, columnId, setCellProps, isExpanded }: EuiDataGridCellValueElementProps) => {
     const { table, alignments, handleFilterClick } = useContext(DataContext);
-    const rawRowValue = table?.rows[rowIndex]?.[columnId];
-    const rowValue = getParsedValue(rawRowValue);
+    const formatter = formatters[columnId];
+    const rawValue: RawValue = table?.rows[rowIndex]?.[columnId];
     const colIndex = columnConfig.columns.findIndex(({ columnId: id }) => id === columnId);
     const {
       oneClickFilter,
@@ -52,13 +41,13 @@ export const createGridCell = (
       colorMapping,
     } = columnConfig.columns[colIndex] ?? {};
     const filterOnClick = oneClickFilter && handleFilterClick;
-    const content = formatters[columnId]?.convert(rawRowValue, filterOnClick ? 'text' : 'html');
+    const content = formatter?.convert(rawValue, filterOnClick ? 'text' : 'html');
     const currentAlignment = alignments?.get(columnId);
 
     useEffect(() => {
       let colorSet = false;
       if (colorMode !== 'none' && (palette || colorMapping)) {
-        const color = getCellColor(columnId, palette, colorMapping)(rowValue);
+        const color = getCellColor(columnId, palette, colorMapping)(rawValue);
 
         if (color) {
           const style = { [colorMode === 'cell' ? 'backgroundColor' : 'color']: color };
@@ -82,7 +71,7 @@ export const createGridCell = (
           });
         };
       }
-    }, [rowValue, columnId, setCellProps, colorMode, palette, colorMapping, isExpanded]);
+    }, [rawValue, columnId, setCellProps, colorMode, palette, colorMapping, isExpanded]);
 
     if (filterOnClick) {
       return (
@@ -95,7 +84,7 @@ export const createGridCell = (
         >
           <EuiLink
             onClick={() => {
-              handleFilterClick?.(columnId, rawRowValue, colIndex, rowIndex);
+              handleFilterClick?.(columnId, rawValue, colIndex, rowIndex);
             }}
           >
             {content}

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

@@ -10,6 +10,7 @@ import { DEFAULT_COLOR_MAPPING_CONFIG } from '@kbn/coloring';
 import { act, screen } from '@testing-library/react';
 import userEvent, { type UserEvent } from '@testing-library/user-event';
 import { chartPluginMock } from '@kbn/charts-plugin/public/mocks';
+import { fieldFormatsServiceMock } from '@kbn/field-formats-plugin/public/mocks';
 import { LayerTypes } from '@kbn/expression-xy-plugin/public';
 import { EuiButtonGroupTestHarness } from '@kbn/test-eui-helpers';
 import { FramePublicAPI, DatasourcePublicAPI, OperationDescriptor, DataType } from '../../../types';
@@ -21,6 +22,8 @@ import { capitalize } from 'lodash';
 import { getKbnPalettes } from '@kbn/palettes';
 import { renderWithProviders } from '../../../test_utils/test_utils';
 
+const fieldFormatsMock = fieldFormatsServiceMock.createStartContract();
+
 describe('data table dimension editor', () => {
   let user: UserEvent;
   let frame: FramePublicAPI;
@@ -94,6 +97,7 @@ describe('data table dimension editor', () => {
       addLayer: jest.fn(),
       removeLayer: jest.fn(),
       datasource: {} as DatasourcePublicAPI,
+      formatFactory: fieldFormatsMock.deserialize,
     };
 
     mockOperationForFirstColumn = (overrides: Partial<OperationDescriptor> = {}) => {

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

@@ -14,6 +14,7 @@ import {
   PaletteOutput,
   PaletteRegistry,
   applyPaletteParams,
+  canCreateCustomMatch,
   getFallbackDataBounds,
 } from '@kbn/coloring';
 import { getColorCategories } from '@kbn/chart-expressions-common';
@@ -33,6 +34,8 @@ import { CollapseSetting } from '../../../shared_components/collapse_setting';
 import { ColorMappingByValues } from '../../../shared_components/coloring/color_mapping_by_values';
 import { ColorMappingByTerms } from '../../../shared_components/coloring/color_mapping_by_terms';
 import { getColumnAlignment } from '../utils';
+import { FormatFactory } from '../../../../common/types';
+import { getDatatableColumn } from '../../../../common/expressions/impl/datatable/utils';
 
 const idPrefix = htmlIdGenerator()();
 
@@ -57,10 +60,11 @@ export type TableDimensionEditorProps =
     paletteService: PaletteRegistry;
     palettes: KbnPalettes;
     isDarkMode: boolean;
+    formatFactory: FormatFactory;
   };
 
 export function TableDimensionEditor(props: TableDimensionEditorProps) {
-  const { frame, accessor, isInlineEditing, isDarkMode } = props;
+  const { frame, accessor, isInlineEditing, isDarkMode, formatFactory } = props;
   const column = props.state.columns.find(({ columnId }) => accessor === columnId);
   const { inputValue: localState, handleInputChange: setLocalState } =
     useDebouncedValue<DatatableVisualizationState>({
@@ -83,6 +87,9 @@ export function TableDimensionEditor(props: TableDimensionEditorProps) {
 
   const currentData =
     frame.activeData?.[localState.layerId] ?? frame.activeData?.[DatatableInspectorTables.Default];
+  const columnMeta = getDatatableColumn(currentData, accessor)?.meta;
+  const formatter = formatFactory(columnMeta?.params);
+  const allowCustomMatch = canCreateCustomMatch(columnMeta);
   const datasource = frame.datasourceLayers?.[localState.layerId];
 
   const { isNumeric, isCategory: isBucketable } = getAccessorType(datasource, accessor);
@@ -109,7 +116,6 @@ 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, [null]);
 
   if (activePalette.name !== CUSTOM_PALETTE && activePalette.params?.stops) {
     activePalette.params.stops = applyPaletteParams(
@@ -247,7 +253,9 @@ export function TableDimensionEditor(props: TableDimensionEditorProps) {
                 }}
                 paletteService={props.paletteService}
                 panelRef={props.panelRef}
-                categories={categories}
+                categories={getColorCategories(currentData?.rows, accessor, [null])}
+                formatter={formatter}
+                allowCustomMatch={allowCustomMatch}
               />
             ) : (
               <ColorMappingByValues

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

@@ -30,13 +30,13 @@ import { CustomPaletteState, EmptyPlaceholder } from '@kbn/charts-plugin/public'
 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 } from '@kbn/transpose-utils';
 import { CoreTheme } from '@kbn/core/public';
 import { getKbnPalettes } from '@kbn/palettes';
+import type { IFieldFormat } from '@kbn/field-formats-plugin/common';
+import { getColorCategories, getLegacyColorCategories } from '@kbn/chart-expressions-common';
 import { css } from '@emotion/react';
 import type { LensTableRowContextMenuEvent } from '../../../types';
-import type { FormatFactory } from '../../../../common/types';
 import { RowHeightMode } from '../../../../common/types';
 import { LensGridDirection } from '../../../../common/expressions';
 import { findMinMaxByColumnId, shouldColorByTerms } from '../../../shared_components';
@@ -61,7 +61,7 @@ import {
 import { getFinalSummaryConfiguration } from '../../../../common/expressions/impl/datatable/summary';
 import { DEFAULT_HEADER_ROW_HEIGHT, DEFAULT_HEADER_ROW_HEIGHT_LINES } from './constants';
 import {
-  getFieldMetaFromDatatable,
+  getDatatableColumn,
   isNumericField,
 } from '../../../../common/expressions/impl/datatable/utils';
 import { CellColorFn, getCellColorFn } from '../../../shared_components/coloring/get_cell_color_fn';
@@ -161,7 +161,7 @@ export const DatatableComponent = (props: DatatableRenderProps) => {
 
   const { getType, dispatchEvent, renderMode, formatFactory, syncColors } = props;
 
-  const formatters: Record<string, ReturnType<FormatFactory>> = useMemo(
+  const formatters: Record<string, IFieldFormat> = useMemo(
     () =>
       firstLocalTable.columns.reduce(
         (map, column) => ({
@@ -401,15 +401,17 @@ export const DatatableComponent = (props: DatatableRenderProps) => {
         return cellColorFnMap.get(originalId)!;
       }
 
-      const dataType = getFieldMetaFromDatatable(firstLocalTable, originalId)?.type;
+      const colInfo = getDatatableColumn(firstLocalTable, originalId);
       const isBucketed = bucketedColumns.some((id) => id === columnId);
-      const colorByTerms = shouldColorByTerms(dataType, isBucketed);
+      const colorByTerms = shouldColorByTerms(colInfo?.meta.type, isBucketed);
       const categoryRows = (untransposedDataRef.current ?? firstLocalTable)?.rows;
+
       const data: ColorMappingInputData = colorByTerms
         ? {
             type: 'categories',
-            // Must use non-transposed data here to correctly collate categories across transposed columns
-            categories: getColorCategories(categoryRows, originalId, [null]),
+            categories: colorMapping
+              ? getColorCategories(categoryRows, originalId, [null])
+              : getLegacyColorCategories(categoryRows, originalId, [null]),
           }
         : {
             type: 'ranges',

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/index.ts

@@ -9,8 +9,8 @@ import type { CoreSetup } from '@kbn/core/public';
 import type { ChartsPluginSetup } from '@kbn/charts-plugin/public';
 import type { ExpressionsSetup } from '@kbn/expressions-plugin/public';
 import type { DataPublicPluginStart } from '@kbn/data-plugin/public';
+import type { FormatFactory } from '@kbn/visualization-ui-components';
 import type { EditorFrameSetup } from '../../types';
-import type { FormatFactory } from '../../../common/types';
 
 interface DatatableVisualizationPluginStartPlugins {
   data: DataPublicPluginStart;

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/runtime_state/converters/index.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { GeneralDatasourceStates } from '../../../../state_management';
+import { convertToRawColorMappingsFn } from './raw_color_mappings';
+
+export const getRuntimeConverters = (datasourceStates?: Readonly<GeneralDatasourceStates>) => [
+  convertToRawColorMappingsFn(datasourceStates),
+];

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/runtime_state/converters/raw_color_mappings.ts

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ColumnState } from '../../../../../common/expressions';
+import {
+  DeprecatedColorMappingConfig,
+  convertToRawColorMappings,
+  getColumnMetaFn,
+  isDeprecatedColorMapping,
+} from '../../../../runtime_state/converters/raw_color_mappings';
+import { GeneralDatasourceStates } from '../../../../state_management';
+import { DatatableVisualizationState } from '../../visualization';
+
+/** @deprecated */
+interface DeprecatedColorMappingColumn extends Omit<ColumnState, 'colorMapping'> {
+  colorMapping: DeprecatedColorMappingConfig;
+}
+
+/**
+ * Old color mapping state meant for type safety during runtime migrations of old configurations
+ *
+ * @deprecated Use respective vis state (i.e. `DatatableVisualizationState`)
+ */
+export interface DeprecatedColorMappingsState extends Omit<DatatableVisualizationState, 'columns'> {
+  columns: Array<DeprecatedColorMappingColumn | ColumnState>;
+}
+
+export const convertToRawColorMappingsFn = (
+  datasourceStates?: Readonly<GeneralDatasourceStates>
+) => {
+  const getColumnMeta = getColumnMetaFn(datasourceStates);
+
+  return (
+    state: DeprecatedColorMappingsState | DatatableVisualizationState
+  ): DatatableVisualizationState => {
+    const hasDeprecatedColorMappings = state.columns.some((column) => {
+      return isDeprecatedColorMapping(column.colorMapping);
+    });
+
+    if (!hasDeprecatedColorMappings) return state as DatatableVisualizationState;
+
+    const convertedColumns = state.columns.map((column) => {
+      if (column.colorMapping?.assignments || column.colorMapping?.specialAssignments) {
+        const columnMeta = getColumnMeta?.(state.layerId, column.columnId);
+
+        return {
+          ...column,
+          colorMapping: convertToRawColorMappings(column.colorMapping, columnMeta),
+        } satisfies ColumnState;
+      }
+
+      return column as ColumnState;
+    });
+
+    return {
+      ...state,
+      columns: convertedColumns,
+    };
+  };
+};

x-pack/platform/plugins/shared/lens/public/visualizations/datatable/runtime_state/index.ts

@@ -0,0 +1,17 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { GeneralDatasourceStates } from '../../../state_management';
+import { DatatableVisualizationState } from '../datatable_visualization';
+import { getRuntimeConverters } from './converters';
+
+export function convertToRuntimeState(
+  state: DatatableVisualizationState,
+  datasourceStates?: Readonly<GeneralDatasourceStates>
+): DatatableVisualizationState {
+  return getRuntimeConverters(datasourceStates).reduce((newState, fn) => fn(newState), state);
+}

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

@@ -6,6 +6,7 @@
  */
 
 import React from 'react';
+
 import { Ast } from '@kbn/interpreter';
 import { i18n } from '@kbn/i18n';
 import { CoreTheme, ThemeServiceStart } from '@kbn/core/public';
@@ -61,6 +62,8 @@ import {
 import { getColorMappingTelemetryEvents } from '../../lens_ui_telemetry/color_telemetry_helpers';
 import { DatatableInspectorTables } from '../../../common/expressions/defs/datatable/datatable';
 import { getSimpleColumnType } from './components/table_actions';
+import { convertToRuntimeState } from './runtime_state';
+
 export interface DatatableVisualizationState {
   columns: ColumnState[];
   layerId: string;
@@ -126,14 +129,18 @@ export const getDatatableVisualization = ({
 
   triggers: [VIS_EVENT_TO_TRIGGER.filter, VIS_EVENT_TO_TRIGGER.tableRowContextMenuClick],
 
-  initialize(addNewLayer, state) {
-    return (
-      state || {
-        columns: [],
-        layerId: addNewLayer(),
-        layerType: LayerTypes.DATA,
-      }
-    );
+  initialize(addNewLayer, state, mainPalette, datasourceStates) {
+    if (state) return convertToRuntimeState(state, datasourceStates);
+
+    return {
+      columns: [],
+      layerId: addNewLayer(),
+      layerType: LayerTypes.DATA,
+    };
+  },
+
+  convertToRuntimeState(state, datasourceStates) {
+    return convertToRuntimeState(state, datasourceStates);
   },
 
   onDatasourceUpdate(state, frame) {
@@ -498,6 +505,7 @@ export const getDatatableVisualization = ({
         isDarkMode={theme.darkMode}
         palettes={palettes}
         paletteService={paletteService}
+        formatFactory={formatFactory}
       />
     );
   },

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

@@ -13,6 +13,7 @@ import { PieVisualizationState } from '../../../common/types';
 import { DimensionEditor, DimensionEditorProps } from './dimension_editor';
 import { getKbnPalettes } from '@kbn/palettes';
 import { chartPluginMock } from '@kbn/charts-plugin/public/mocks';
+import { fieldFormatsServiceMock } from '@kbn/field-formats-plugin/public/mocks';
 
 const darkMode = false;
 const paletteServiceMock = chartPluginMock.createPaletteRegistry();
@@ -39,7 +40,7 @@ describe('DimensionEditor', () => {
           colorMapping: {
             assignments: [],
             specialAssignments: [
-              { rule: { type: 'other' }, color: { type: 'loop' }, touched: false },
+              { rules: [{ type: 'other' }], color: { type: 'loop' }, touched: false },
             ],
             paletteId: 'default',
             colorMode: { type: 'categorical' },
@@ -49,6 +50,7 @@ describe('DimensionEditor', () => {
     };
 
     const mockFrame = createMockFramePublicAPI();
+    const fieldFormatsMock = fieldFormatsServiceMock.createStartContract();
     mockFrame.datasourceLayers = Object.fromEntries(
       defaultState.layers.map(({ layerId: id }) => [id, createMockDatasource(id).publicAPIMock])
     );
@@ -67,6 +69,7 @@ describe('DimensionEditor', () => {
       palettes,
       isDarkMode: darkMode,
       paletteService: paletteServiceMock,
+      formatFactory: fieldFormatsMock.deserialize,
     };
 
     buildProps = (props = {}) => {

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

@@ -14,7 +14,7 @@ import {
   ColorMapping,
   SPECIAL_TOKENS_STRING_CONVERSION,
 } from '@kbn/coloring';
-import { ColorPicker } from '@kbn/visualization-ui-components';
+import { ColorPicker, FormatFactory } from '@kbn/visualization-ui-components';
 import { useDebouncedValue } from '@kbn/visualization-utils';
 import { EuiFormRow, EuiFlexGroup, EuiFlexItem, EuiSwitch, EuiText, EuiBadge } from '@elastic/eui';
 import { useState, useCallback } from 'react';
@@ -34,9 +34,11 @@ import {
   isCollapsed,
 } from './visualization';
 import { trackUiCounterEvents } from '../../lens_ui_telemetry';
+import { getDatatableColumn } from '../../../common/expressions/impl/datatable/utils';
 import { getSortedAccessorsForGroup } from './to_expression';
 
 export type DimensionEditorProps = VisualizationDimensionEditorProps<PieVisualizationState> & {
+  formatFactory: FormatFactory;
   paletteService: PaletteRegistry;
   palettes: KbnPalettes;
   isDarkMode: boolean;
@@ -50,9 +52,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
     });
 
   const currentLayer = localState.layers.find((layer) => layer.layerId === props.layerId);
-
-  const canUseColorMapping = currentLayer && currentLayer.colorMapping ? true : false;
-  const [useNewColorMapping, setUseNewColorMapping] = useState(canUseColorMapping);
+  const [useNewColorMapping, setUseNewColorMapping] = useState(Boolean(currentLayer?.colorMapping));
 
   const setConfig = useCallback(
     ({ color }: { color?: string }) => {
@@ -132,8 +132,9 @@ export function DimensionEditor(props: DimensionEditorProps) {
     props.state.palette,
     currentLayer.colorMapping
   );
-  const table = props.frame.activeData?.[currentLayer.layerId];
-  const splitCategories = getColorCategories(table?.rows, props.accessor);
+  const currentData = props.frame.activeData?.[currentLayer.layerId];
+  const columnMeta = getDatatableColumn(currentData, props.accessor)?.meta;
+  const formatter = props.formatFactory(columnMeta?.params);
 
   return (
     <>
@@ -143,7 +144,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
           label={i18n.translate('xpack.lens.colorMapping.editColorMappingSectionLabel', {
             defaultMessage: 'Color mapping',
           })}
-          style={{ alignItems: 'center' }}
+          css={{ alignItems: 'center' }}
           fullWidth
         >
           <PalettePanelContainer
@@ -191,7 +192,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
                   />
                 </EuiFlexItem>
                 <EuiFlexItem>
-                  {canUseColorMapping || useNewColorMapping ? (
+                  {useNewColorMapping ? (
                     <CategoricalColorMapping
                       isDarkMode={props.isDarkMode}
                       model={currentLayer.colorMapping ?? { ...DEFAULT_COLOR_MAPPING_CONFIG }}
@@ -199,9 +200,10 @@ export function DimensionEditor(props: DimensionEditorProps) {
                       palettes={props.palettes}
                       data={{
                         type: 'categories',
-                        categories: splitCategories,
+                        categories: getColorCategories(currentData?.rows, props.accessor),
                       }}
                       specialTokens={SPECIAL_TOKENS_STRING_CONVERSION}
+                      formatter={formatter}
                     />
                   ) : (
                     <PalettePicker

x-pack/platform/plugins/shared/lens/public/visualizations/partition/index.ts

@@ -8,11 +8,13 @@
 import type { CoreSetup } from '@kbn/core/public';
 import type { UiActionsStart } from '@kbn/ui-actions-plugin/public';
 import type { ChartsPluginSetup } from '@kbn/charts-plugin/public';
+import { FormatFactory } from '@kbn/visualization-ui-components';
 import type { EditorFrameSetup } from '../../types';
 
 export interface PieVisualizationPluginSetupPlugins {
   editorFrame: EditorFrameSetup;
   charts: ChartsPluginSetup;
+  formatFactory: FormatFactory;
 }
 
 export interface PieVisualizationPluginStartPlugins {
@@ -20,14 +22,17 @@ export interface PieVisualizationPluginStartPlugins {
 }
 
 export class PieVisualization {
-  setup(core: CoreSetup, { editorFrame, charts }: PieVisualizationPluginSetupPlugins) {
+  setup(
+    core: CoreSetup,
+    { editorFrame, formatFactory, charts }: PieVisualizationPluginSetupPlugins
+  ) {
     editorFrame.registerVisualization(async () => {
       const [{ getPieVisualization }, paletteService] = await Promise.all([
         import('../../async_services'),
         charts.palettes.getPalettes(),
       ]);
 
-      return getPieVisualization({ paletteService, kibanaTheme: core.theme });
+      return getPieVisualization({ paletteService, kibanaTheme: core.theme, formatFactory });
     });
   }
 }

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

@@ -1,40 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-
-import { LegendValue } from '@elastic/charts';
-import { cloneDeep } from 'lodash';
-import { PieLayerState, PieVisualizationState } from '../../../common/types';
-
-type PersistedPieLayerState = PieLayerState & {
-  showValuesInLegend?: boolean;
-};
-
-export type PersistedPieVisualizationState = Omit<PieVisualizationState, 'layers'> & {
-  layers: PersistedPieLayerState[];
-};
-
-export function convertToRuntime(state: PersistedPieVisualizationState) {
-  let newState = cloneDeep(state) as unknown as PieVisualizationState;
-  newState = convertToLegendStats(newState);
-  return newState;
-}
-
-function convertToLegendStats(state: PieVisualizationState) {
-  state.layers.forEach((l) => {
-    if ('showValuesInLegend' in l) {
-      l.legendStats = [
-        ...new Set([
-          ...(l.showValuesInLegend ? [LegendValue.Value] : []),
-          ...(l.legendStats || []),
-        ]),
-      ];
-    }
-    delete (l as PersistedPieLayerState).showValuesInLegend;
-  });
-
-  return state;
-}