[ES|QL] [Discover] Keeps the preferred chart configuration when possible (#197453)

## Summary

Closes https://github.com/elastic/kibana/issues/184631

It keeps the chart configuration when the user is doing actions
compatible with the current query such as:

- Adding a where filter (by clicking the table, the sidebar, the chart)
- Changes the breakdown field and the field type is compatible with the
current chart
- Changing to a compatible chart type (from example from bar to line or
pie to treemap)
- Changing the query that doesnt affect the generated columns mapped to
a chart. For example adding a limit or creating a runtime field etc.

The logic depends on the suggestions. If the suggestions return the
preferred chart type, then we are going to use this. So it really
depends on the api and the type / number of columns. It is as smarter as
it can in order to not create bugs. I am quite happy with the result. It
is much better than what we have so far.


![meow](https://github.com/user-attachments/assets/c4249e5e-e785-4e57-8651-d1f660f5a61a)

### Next steps
I would love to do the same on the dahsboard too, needs more time
though. But the changes made here will def work in favor

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [x] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed

---------

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

packages/kbn-visualization-utils/index.ts

@@ -11,3 +11,6 @@ export { getTimeZone } from './src/get_timezone';
 export { getLensAttributesFromSuggestion } from './src/get_lens_attributes';
 export { TooltipWrapper } from './src/tooltip_wrapper';
 export { useDebouncedValue } from './src/debounced_value';
+export { ChartType } from './src/types';
+export { getDatasourceId } from './src/get_datasource_id';
+export { mapVisToChartType } from './src/map_vis_to_chart_type';

packages/kbn-visualization-utils/src/get_datasource_id.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", 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".
+ */
+
+export const getDatasourceId = (datasourceStates: Record<string, unknown>) => {
+  const datasourceId: 'formBased' | 'textBased' | undefined = [
+    'formBased' as const,
+    'textBased' as const,
+  ].find((key) => Boolean(datasourceStates[key]));
+
+  return datasourceId;
+};

packages/kbn-visualization-utils/src/map_vis_to_chart_type.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 { ChartType, LensVisualizationType } from './types';
+
+type ValueOf<T> = T[keyof T];
+type LensToChartMap = {
+  [K in ValueOf<typeof LensVisualizationType>]: ChartType;
+};
+const lensTypesToChartTypes: LensToChartMap = {
+  [LensVisualizationType.XY]: ChartType.XY,
+  [LensVisualizationType.Metric]: ChartType.Metric,
+  [LensVisualizationType.LegacyMetric]: ChartType.Metric,
+  [LensVisualizationType.Pie]: ChartType.Pie,
+  [LensVisualizationType.Heatmap]: ChartType.Heatmap,
+  [LensVisualizationType.Gauge]: ChartType.Gauge,
+  [LensVisualizationType.Datatable]: ChartType.Table,
+};
+function isLensVisualizationType(value: string): value is LensVisualizationType {
+  return Object.values(LensVisualizationType).includes(value as LensVisualizationType);
+}
+export const mapVisToChartType = (visualizationType: string) => {
+  if (isLensVisualizationType(visualizationType)) {
+    return lensTypesToChartTypes[visualizationType];
+  }
+};

packages/kbn-visualization-utils/src/types.ts

@@ -43,3 +43,30 @@ export interface Suggestion<T = unknown, V = unknown> {
   changeType: TableChangeType;
   keptLayerIds: string[];
 }
+
+export enum ChartType {
+  XY = 'XY',
+  Gauge = 'Gauge',
+  Bar = 'Bar',
+  Line = 'Line',
+  Area = 'Area',
+  Donut = 'Donut',
+  Heatmap = 'Heatmap',
+  Metric = 'Metric',
+  Treemap = 'Treemap',
+  Tagcloud = 'Tagcloud',
+  Waffle = 'Waffle',
+  Pie = 'Pie',
+  Mosaic = 'Mosaic',
+  Table = 'Table',
+}
+
+export enum LensVisualizationType {
+  XY = 'lnsXY',
+  Metric = 'lnsMetric',
+  Pie = 'lnsPie',
+  Heatmap = 'lnsHeatmap',
+  Gauge = 'lnsGauge',
+  Datatable = 'lnsDatatable',
+  LegacyMetric = 'lnsLegacyMetric',
+}

src/plugins/unified_histogram/public/services/lens_vis_service.ts

@@ -27,7 +27,11 @@ import type {
 import type { AggregateQuery, TimeRange } from '@kbn/es-query';
 import { getAggregateQueryMode, isOfAggregateQueryType } from '@kbn/es-query';
 import { i18n } from '@kbn/i18n';
-import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import {
+  getLensAttributesFromSuggestion,
+  ChartType,
+  mapVisToChartType,
+} from '@kbn/visualization-utils';
 import { LegendSize } from '@kbn/visualizations-plugin/public';
 import { XYConfiguration } from '@kbn/visualizations-plugin/common';
 import type { Datatable, DatatableColumn } from '@kbn/expressions-plugin/common';
@@ -42,6 +46,7 @@ import {
   isSuggestionShapeAndVisContextCompatible,
   deriveLensSuggestionFromLensAttributes,
   type QueryParams,
+  injectESQLQueryIntoLensLayers,
 } from '../utils/external_vis_context';
 import { computeInterval } from '../utils/compute_interval';
 import { fieldSupportsBreakdown } from '../utils/field_supports_breakdown';
@@ -147,7 +152,10 @@ export class LensVisService {
       externalVisContextStatus: UnifiedHistogramExternalVisContextStatus
     ) => void;
   }) => {
-    const allSuggestions = this.getAllSuggestions({ queryParams });
+    const allSuggestions = this.getAllSuggestions({
+      queryParams,
+      preferredVisAttributes: externalVisContext?.attributes,
+    });
 
     const suggestionState = this.getCurrentSuggestionState({
       externalVisContext,
@@ -252,6 +260,7 @@ export class LensVisService {
       const histogramSuggestionForESQL = this.getHistogramSuggestionForESQL({
         queryParams,
         breakdownField,
+        preferredVisAttributes: externalVisContext?.attributes,
       });
       if (histogramSuggestionForESQL) {
         // In case if histogram suggestion, we want to empty the array and push the new suggestion
@@ -463,9 +472,11 @@ export class LensVisService {
   private getHistogramSuggestionForESQL = ({
     queryParams,
     breakdownField,
+    preferredVisAttributes,
   }: {
     queryParams: QueryParams;
     breakdownField?: DataViewField;
+    preferredVisAttributes?: UnifiedHistogramVisContext['attributes'];
   }): Suggestion | undefined => {
     const { dataView, query, timeRange, columns } = queryParams;
     const breakdownColumn = breakdownField?.name
@@ -510,7 +521,22 @@ export class LensVisService {
       if (breakdownColumn) {
         context.textBasedColumns.push(breakdownColumn);
       }
-      const suggestions = this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? [];
+
+      // here the attributes contain the main query and not the histogram one
+      const updatedAttributesWithQuery = preferredVisAttributes
+        ? injectESQLQueryIntoLensLayers(preferredVisAttributes, {
+            esql: esqlQuery,
+          })
+        : undefined;
+
+      const suggestions =
+        this.lensSuggestionsApi(
+          context,
+          dataView,
+          ['lnsDatatable'],
+          ChartType.XY,
+          updatedAttributesWithQuery
+        ) ?? [];
       if (suggestions.length) {
         const suggestion = suggestions[0];
         const suggestionVisualizationState = Object.assign({}, suggestion?.visualizationState);
@@ -574,9 +600,25 @@ export class LensVisService {
     );
   };
 
-  private getAllSuggestions = ({ queryParams }: { queryParams: QueryParams }): Suggestion[] => {
+  private getAllSuggestions = ({
+    queryParams,
+    preferredVisAttributes,
+  }: {
+    queryParams: QueryParams;
+    preferredVisAttributes?: UnifiedHistogramVisContext['attributes'];
+  }): Suggestion[] => {
     const { dataView, columns, query, isPlainRecord } = queryParams;
 
+    const preferredChartType = preferredVisAttributes
+      ? mapVisToChartType(preferredVisAttributes.visualizationType)
+      : undefined;
+
+    let visAttributes = preferredVisAttributes;
+
+    if (query && isOfAggregateQueryType(query) && preferredVisAttributes) {
+      visAttributes = injectESQLQueryIntoLensLayers(preferredVisAttributes, query);
+    }
+
     const context = {
       dataViewSpec: dataView?.toSpec(),
       fieldName: '',
@@ -584,7 +626,13 @@ export class LensVisService {
       query: query && isOfAggregateQueryType(query) ? query : undefined,
     };
     const allSuggestions = isPlainRecord
-      ? this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? []
+      ? this.lensSuggestionsApi(
+          context,
+          dataView,
+          ['lnsDatatable'],
+          preferredChartType,
+          visAttributes
+        ) ?? []
       : [];
 
     return allSuggestions;

src/plugins/unified_histogram/public/utils/external_vis_context.test.ts

@@ -13,6 +13,7 @@ import {
   canImportVisContext,
   exportVisContext,
   isSuggestionShapeAndVisContextCompatible,
+  injectESQLQueryIntoLensLayers,
 } from './external_vis_context';
 import { getLensVisMock } from '../__mocks__/lens_vis';
 import { dataViewWithTimefieldMock } from '../__mocks__/data_view_with_timefield';
@@ -162,4 +163,63 @@ describe('external_vis_context', () => {
       ).toBe(true);
     });
   });
+
+  describe('injectESQLQueryIntoLensLayers', () => {
+    it('should return the Lens attributes as they are for unknown datasourceId', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { unknownId: { layers: {} } },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+      expect(injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo' })).toStrictEqual(
+        attributes
+      );
+    });
+
+    it('should return the Lens attributes as they are for DSL config (formbased)', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { formBased: { layers: {} } },
+        },
+      } as UnifiedHistogramVisContext['attributes'];
+      expect(injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo' })).toStrictEqual(
+        attributes
+      );
+    });
+
+    it('should inject the query to the Lens attributes for ES|QL config (textbased)', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { textBased: { layers: { layer1: { query: { esql: 'from foo' } } } } },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+
+      const expectedAttributes = {
+        ...attributes,
+        state: {
+          ...attributes.state,
+          datasourceStates: {
+            ...attributes.state.datasourceStates,
+            textBased: {
+              ...attributes.state.datasourceStates.textBased,
+              layers: {
+                layer1: {
+                  query: { esql: 'from foo | stats count(*)' },
+                },
+              },
+            },
+          },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+      expect(
+        injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo | stats count(*)' })
+      ).toStrictEqual(expectedAttributes);
+    });
+  });
 });

src/plugins/unified_histogram/public/utils/external_vis_context.ts

@@ -7,9 +7,10 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { isEqual } from 'lodash';
+import { isEqual, cloneDeep } from 'lodash';
 import type { DataView } from '@kbn/data-views-plugin/common';
 import type { AggregateQuery, Filter, Query, TimeRange } from '@kbn/es-query';
+import { getDatasourceId } from '@kbn/visualization-utils';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
 import type { PieVisualizationState, Suggestion, XYState } from '@kbn/lens-plugin/public';
 import { UnifiedHistogramSuggestionType, UnifiedHistogramVisContext } from '../types';
@@ -103,6 +104,42 @@ export const isSuggestionShapeAndVisContextCompatible = (
   );
 };
 
+export const injectESQLQueryIntoLensLayers = (
+  visAttributes: UnifiedHistogramVisContext['attributes'],
+  query: AggregateQuery
+) => {
+  const datasourceId = getDatasourceId(visAttributes.state.datasourceStates);
+
+  // if the datasource is formBased, we should not fix the query
+  if (!datasourceId || datasourceId === 'formBased') {
+    return visAttributes;
+  }
+
+  if (!visAttributes.state.datasourceStates[datasourceId]) {
+    return visAttributes;
+  }
+
+  const datasourceState = cloneDeep(visAttributes.state.datasourceStates[datasourceId]);
+
+  if (datasourceState && datasourceState.layers) {
+    Object.values(datasourceState.layers).forEach((layer) => {
+      if (!isEqual(layer.query, query)) {
+        layer.query = query;
+      }
+    });
+  }
+  return {
+    ...visAttributes,
+    state: {
+      ...visAttributes.state,
+      datasourceStates: {
+        ...visAttributes.state.datasourceStates,
+        [datasourceId]: datasourceState,
+      },
+    },
+  };
+};
+
 export function deriveLensSuggestionFromLensAttributes({
   externalVisContext,
   queryParams,
@@ -122,10 +159,7 @@ export function deriveLensSuggestionFromLensAttributes({
       }
 
       // it should be one of 'formBased'/'textBased' and have value
-      const datasourceId: 'formBased' | 'textBased' | undefined = [
-        'formBased' as const,
-        'textBased' as const,
-      ].find((key) => Boolean(externalVisContext.attributes.state.datasourceStates[key]));
+      const datasourceId = getDatasourceId(externalVisContext.attributes.state.datasourceStates);
 
       if (!datasourceId) {
         return undefined;

test/functional/apps/discover/esql/_esql_view.ts

@@ -676,6 +676,94 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
           `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB | where countB > 0\nAND \`geo.dest\`=="BT"`
         );
       });
+
+      it('should append a where clause by clicking the table without changing the chart type', async () => {
+        await discover.selectTextBaseLang();
+        const testQuery = `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB`;
+        await monacoEditor.setCodeEditorValue(testQuery);
+
+        await testSubjects.click('querySubmitButton');
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        // change the type to line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lnsChartSwitchPopover');
+        await testSubjects.click('lnsChartSwitchPopover_line');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('applyFlyoutButton');
+
+        await dataGrid.clickCellFilterForButtonExcludingControlColumns(0, 1);
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        const editorValue = await monacoEditor.getCodeEditorValue();
+        expect(editorValue).to.eql(
+          `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB\n| WHERE \`geo.dest\`=="BT"`
+        );
+
+        // check that the type is still line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        const chartSwitcher = await testSubjects.find('lnsChartSwitchPopover');
+        const type = await chartSwitcher.getVisibleText();
+        expect(type).to.be('Line');
+      });
+
+      it('should append a where clause by clicking the table without changing the chart type nor the visualization state', async () => {
+        await discover.selectTextBaseLang();
+        const testQuery = `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB`;
+        await monacoEditor.setCodeEditorValue(testQuery);
+
+        await testSubjects.click('querySubmitButton');
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        // change the type to line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lnsChartSwitchPopover');
+        await testSubjects.click('lnsChartSwitchPopover_line');
+
+        // change the color to red
+        await testSubjects.click('lnsXY_yDimensionPanel');
+        const colorPickerInput = await testSubjects.find('~indexPattern-dimension-colorPicker');
+        await colorPickerInput.clearValueWithKeyboard();
+        await colorPickerInput.type('#ff0000');
+        await common.sleep(1000); // give time for debounced components to rerender
+
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lns-indexPattern-dimensionContainerClose');
+        await testSubjects.click('applyFlyoutButton');
+
+        await dataGrid.clickCellFilterForButtonExcludingControlColumns(0, 1);
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        const editorValue = await monacoEditor.getCodeEditorValue();
+        expect(editorValue).to.eql(
+          `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB\n| WHERE \`geo.dest\`=="BT"`
+        );
+
+        // check that the type is still line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        const chartSwitcher = await testSubjects.find('lnsChartSwitchPopover');
+        const type = await chartSwitcher.getVisibleText();
+        expect(type).to.be('Line');
+
+        // check that the color is still red
+        await testSubjects.click('lnsXY_yDimensionPanel');
+        const colorPickerInputAfterFilter = await testSubjects.find(
+          '~indexPattern-dimension-colorPicker'
+        );
+        expect(await colorPickerInputAfterFilter.getAttribute('value')).to.be('#FF0000');
+      });
     });
 
     describe('histogram breakdown', () => {

test/functional/apps/discover/group3/_lens_vis.ts

@@ -288,7 +288,8 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       await header.waitUntilLoadingHasFinished();
       await discover.waitUntilSearchingHasFinished();
 
-      expect(await getCurrentVisTitle()).to.be('Bar');
+      // Line has been retained although the query changed!
+      expect(await getCurrentVisTitle()).to.be('Line');
 
       await checkESQLHistogramVis(defaultTimespanESQL, '100');
 
@@ -567,15 +568,14 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       await testSubjects.existOrFail('partitionVisChart');
       expect(await discover.getVisContextSuggestionType()).to.be('lensSuggestion');
 
-      await monacoEditor.setCodeEditorValue(
-        'from logstash-* | stats averageB = avg(bytes) by extension.raw'
-      );
+      // reset to histogram
+      await monacoEditor.setCodeEditorValue('from logstash-*');
       await testSubjects.click('querySubmitButton');
       await header.waitUntilLoadingHasFinished();
       await discover.waitUntilSearchingHasFinished();
 
       expect(await getCurrentVisTitle()).to.be('Bar');
-      expect(await discover.getVisContextSuggestionType()).to.be('lensSuggestion');
+      expect(await discover.getVisContextSuggestionType()).to.be('histogramForESQL');
 
       await testSubjects.existOrFail('unsavedChangesBadge');
 

x-pack/plugins/lens/public/lens_suggestions_api/helpers.test.ts

@@ -0,0 +1,206 @@
+/*
+ * 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 type { DatatableColumn } from '@kbn/expressions-plugin/common';
+import { mergeSuggestionWithVisContext } from './helpers';
+import { mockAllSuggestions } from '../mocks';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+
+const context = {
+  dataViewSpec: {
+    id: 'index1',
+    title: 'index1',
+    name: 'DataView',
+  },
+  fieldName: '',
+  textBasedColumns: [
+    {
+      id: 'field1',
+      name: 'field1',
+      meta: {
+        type: 'number',
+      },
+    },
+    {
+      id: 'field2',
+      name: 'field2',
+      meta: {
+        type: 'string',
+      },
+    },
+  ] as DatatableColumn[],
+  query: {
+    esql: 'FROM index1 | keep field1, field2',
+  },
+};
+
+describe('lens suggestions api helpers', () => {
+  describe('mergeSuggestionWithVisContext', () => {
+    it('should return the suggestion as it is if the visualization types do not match', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion as it is if the context is not from ES|QL', async () => {
+      const nonESQLContext = {
+        dataViewSpec: {
+          id: 'index1',
+          title: 'index1',
+          name: 'DataView',
+        },
+        fieldName: 'field1',
+      };
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(
+        mergeSuggestionWithVisContext({ suggestion, visAttributes, context: nonESQLContext })
+      ).toStrictEqual(suggestion);
+    });
+
+    it('should return the suggestion as it is for DSL config (formbased)', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { formBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion as it is for columns that dont match the context', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            shape: 'heatmap',
+          },
+          datasourceStates: {
+            textBased: {
+              layers: {
+                layer1: {
+                  index: 'layer1',
+                  query: {
+                    esql: 'FROM kibana_sample_data_flights | keep Dest, AvgTicketPrice',
+                  },
+                  columns: [
+                    {
+                      columnId: 'colA',
+                      fieldName: 'Dest',
+                      meta: {
+                        type: 'string',
+                      },
+                    },
+                    {
+                      columnId: 'colB',
+                      fieldName: 'AvgTicketPrice',
+                      meta: {
+                        type: 'number',
+                      },
+                    },
+                  ],
+                  timeField: 'timestamp',
+                },
+              },
+            },
+          },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion updated with the attributes if the visualization types and the context columns match', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            shape: 'heatmap',
+            layerId: 'layer1',
+            layerType: 'data',
+            legend: {
+              isVisible: false,
+              position: 'left',
+              type: 'heatmap_legend',
+            },
+            gridConfig: {
+              type: 'heatmap_grid',
+              isCellLabelVisible: true,
+              isYAxisLabelVisible: false,
+              isXAxisLabelVisible: false,
+              isYAxisTitleVisible: false,
+              isXAxisTitleVisible: false,
+            },
+            valueAccessor: 'acc1',
+            xAccessor: 'acc2',
+          },
+          datasourceStates: {
+            textBased: {
+              layers: {
+                layer1: {
+                  index: 'layer1',
+                  query: {
+                    esql: 'FROM index1 | keep field1, field2',
+                  },
+                  columns: [
+                    {
+                      columnId: 'field2',
+                      fieldName: 'field2',
+                      meta: {
+                        type: 'string',
+                      },
+                    },
+                    {
+                      columnId: 'field1',
+                      fieldName: 'field1',
+                      meta: {
+                        type: 'number',
+                      },
+                    },
+                  ],
+                  timeField: 'timestamp',
+                },
+              },
+            },
+          },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      const updatedSuggestion = mergeSuggestionWithVisContext({
+        suggestion,
+        visAttributes,
+        context,
+      });
+      expect(updatedSuggestion.visualizationState).toStrictEqual(visAttributes.state.visualization);
+    });
+  });
+});

x-pack/plugins/lens/public/lens_suggestions_api/helpers.ts

@@ -0,0 +1,76 @@
+/*
+ * 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 type { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
+import { getDatasourceId } from '@kbn/visualization-utils';
+import type { VisualizeEditorContext, Suggestion } from '../types';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+
+/**
+ * Returns the suggestion updated with external visualization state for ES|QL charts
+ * The visualization state is merged with the suggestion if the datasource is textBased, the columns match the context and the visualization type matches
+ * @param suggestion the suggestion to be updated
+ * @param visAttributes the preferred visualization attributes
+ * @param context the lens suggestions api context as being set by the consumers
+ * @returns updated suggestion
+ */
+
+export function mergeSuggestionWithVisContext({
+  suggestion,
+  visAttributes,
+  context,
+}: {
+  suggestion: Suggestion;
+  visAttributes: TypedLensByValueInput['attributes'];
+  context: VisualizeFieldContext | VisualizeEditorContext;
+}): Suggestion {
+  if (
+    visAttributes.visualizationType !== suggestion.visualizationId ||
+    !('textBasedColumns' in context)
+  ) {
+    return suggestion;
+  }
+
+  // it should be one of 'formBased'/'textBased' and have value
+  const datasourceId = getDatasourceId(visAttributes.state.datasourceStates);
+
+  // if the datasource is formBased, we should not merge
+  if (!datasourceId || datasourceId === 'formBased') {
+    return suggestion;
+  }
+  const datasourceState = Object.assign({}, visAttributes.state.datasourceStates[datasourceId]);
+
+  // should be based on same columns
+  if (
+    !datasourceState?.layers ||
+    Object.values(datasourceState?.layers).some(
+      (layer) =>
+        layer.columns?.some(
+          (c: { fieldName: string }) =>
+            !context?.textBasedColumns?.find((col) => col.name === c.fieldName)
+        ) || layer.columns?.length !== context?.textBasedColumns?.length
+    )
+  ) {
+    return suggestion;
+  }
+  const layerIds = Object.keys(datasourceState.layers);
+  try {
+    return {
+      title: visAttributes.title,
+      visualizationId: visAttributes.visualizationType,
+      visualizationState: visAttributes.state.visualization,
+      keptLayerIds: layerIds,
+      datasourceState,
+      datasourceId,
+      columns: suggestion.columns,
+      changeType: suggestion.changeType,
+      score: suggestion.score,
+      previewIcon: suggestion.previewIcon,
+    };
+  } catch {
+    return suggestion;
+  }
+}

x-pack/plugins/lens/public/lens_suggestions_api/index.ts

@@ -6,22 +6,12 @@
  */
 import type { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
 import type { DataView } from '@kbn/data-views-plugin/public';
-import { getSuggestions } from './editor_frame_service/editor_frame/suggestion_helpers';
-import type { DatasourceMap, VisualizationMap, VisualizeEditorContext } from './types';
-import type { DataViewsState } from './state_management';
-
-export enum ChartType {
-  XY = 'XY',
-  Bar = 'Bar',
-  Line = 'Line',
-  Area = 'Area',
-  Donut = 'Donut',
-  Heatmap = 'Heat map',
-  Treemap = 'Treemap',
-  Tagcloud = 'Tag cloud',
-  Waffle = 'Waffle',
-  Table = 'Table',
-}
+import type { ChartType } from '@kbn/visualization-utils';
+import { getSuggestions } from '../editor_frame_service/editor_frame/suggestion_helpers';
+import type { DatasourceMap, VisualizationMap, VisualizeEditorContext } from '../types';
+import type { DataViewsState } from '../state_management';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+import { mergeSuggestionWithVisContext } from './helpers';
 
 interface SuggestionsApiProps {
   context: VisualizeFieldContext | VisualizeEditorContext;
@@ -30,6 +20,7 @@ interface SuggestionsApiProps {
   datasourceMap?: DatasourceMap;
   excludedVisualizations?: string[];
   preferredChartType?: ChartType;
+  preferredVisAttributes?: TypedLensByValueInput['attributes'];
 }
 
 export const suggestionsApi = ({
@@ -39,6 +30,7 @@ export const suggestionsApi = ({
   visualizationMap,
   excludedVisualizations,
   preferredChartType,
+  preferredVisAttributes,
 }: SuggestionsApiProps) => {
   const initialContext = context;
   if (!datasourceMap || !visualizationMap || !dataView.id) return undefined;
@@ -79,32 +71,7 @@ export const suggestionsApi = ({
     dataViews,
   });
   if (!suggestions.length) return [];
-  // check if there is an XY chart suggested
-  // if user has requested for a line or area, we want to sligthly change the state
-  // to return line / area instead of a bar chart
-  const chartType = preferredChartType?.toLowerCase();
-  const XYSuggestion = suggestions.find((sug) => sug.visualizationId === 'lnsXY');
-  if (XYSuggestion && chartType && ['area', 'line'].includes(chartType)) {
-    const visualizationState = visualizationMap[
-      XYSuggestion.visualizationId
-    ]?.switchVisualizationType?.(chartType, XYSuggestion?.visualizationState);
-    return [
-      {
-        ...XYSuggestion,
-        visualizationState,
-      },
-    ];
-  }
-  // in case the user asks for another type (except from area, line) check if it exists
-  // in suggestions and return this instead
-  if (suggestions.length > 1 && preferredChartType) {
-    const suggestionFromModel = suggestions.find(
-      (s) => s.title.includes(preferredChartType) || s.visualizationId.includes(preferredChartType)
-    );
-    if (suggestionFromModel) {
-      return [suggestionFromModel];
-    }
-  }
+
   const activeVisualization = suggestions[0];
   if (
     activeVisualization.incomplete ||
@@ -126,7 +93,46 @@ export const suggestionsApi = ({
     visualizationState: activeVisualization.visualizationState,
     dataViews,
   }).filter((sug) => !sug.hide && sug.visualizationId !== 'lnsLegacyMetric');
+
+  // check if there is an XY chart suggested
+  // if user has requested for a line or area, we want to sligthly change the state
+  // to return line / area instead of a bar chart
+  const chartType = preferredChartType?.toLowerCase();
+  const XYSuggestion = newSuggestions.find((s) => s.visualizationId === 'lnsXY');
+  // a type can be area, line, area_stacked, area_percentage etc
+  const isAreaOrLine = ['area', 'line'].some((type) => chartType?.includes(type));
+  if (XYSuggestion && chartType && isAreaOrLine) {
+    const visualizationState = visualizationMap[
+      XYSuggestion.visualizationId
+    ]?.switchVisualizationType?.(chartType, XYSuggestion?.visualizationState);
+
+    return [
+      {
+        ...XYSuggestion,
+        visualizationState,
+      },
+    ];
+  }
+  // in case the user asks for another type (except from area, line) check if it exists
+  // in suggestions and return this instead
   const suggestionsList = [activeVisualization, ...newSuggestions];
+  if (suggestionsList.length > 1 && preferredChartType) {
+    const compatibleSuggestion = suggestionsList.find(
+      (s) => s.title.includes(preferredChartType) || s.visualizationId.includes(preferredChartType)
+    );
+
+    if (compatibleSuggestion) {
+      const suggestion = preferredVisAttributes
+        ? mergeSuggestionWithVisContext({
+            suggestion: compatibleSuggestion,
+            visAttributes: preferredVisAttributes,
+            context,
+          })
+        : compatibleSuggestion;
+
+      return [suggestion];
+    }
+  }
 
   // if there is no preference from the user, send everything
   // until we separate the text based suggestions logic from the dataview one,

x-pack/plugins/lens/public/lens_suggestions_api/lens_suggestions_api.test.ts

@@ -6,9 +6,11 @@
  */
 import type { DataView } from '@kbn/data-views-plugin/public';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
-import { createMockVisualization, DatasourceMock, createMockDatasource } from './mocks';
-import { DatasourceSuggestion } from './types';
-import { suggestionsApi, ChartType } from './lens_suggestions_api';
+import { ChartType } from '@kbn/visualization-utils';
+import { createMockVisualization, DatasourceMock, createMockDatasource } from '../mocks';
+import { DatasourceSuggestion } from '../types';
+import { suggestionsApi } from '.';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
 
 const generateSuggestion = (state = {}, layerId: string = 'first'): DatasourceSuggestion => ({
   state,
@@ -264,6 +266,9 @@ describe('suggestionsApi', () => {
     datasourceMap.textBased.getDatasourceSuggestionsForVisualizeField.mockReturnValue([
       generateSuggestion(),
     ]);
+    datasourceMap.textBased.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      generateSuggestion(),
+    ]);
     const context = {
       dataViewSpec: {
         id: 'index1',
@@ -284,8 +289,7 @@ describe('suggestionsApi', () => {
       preferredChartType: ChartType.Line,
     });
     expect(suggestions?.length).toEqual(1);
-    expect(suggestions?.[0]).toMatchInlineSnapshot(
-      `
+    expect(suggestions?.[0]).toMatchInlineSnapshot(`
       Object {
         "changeType": "unchanged",
         "columns": 0,
@@ -302,8 +306,111 @@ describe('suggestionsApi', () => {
           "preferredSeriesType": "line",
         },
       }
-    `
-    );
+    `);
+  });
+
+  test('returns the suggestion with the preferred attributes ', async () => {
+    const dataView = { id: 'index1' } as unknown as DataView;
+    const visualizationMap = {
+      lnsXY: {
+        ...mockVis,
+        switchVisualizationType(seriesType: string, state: unknown) {
+          return {
+            ...(state as Record<string, unknown>),
+            preferredSeriesType: seriesType,
+          };
+        },
+        getSuggestions: () => [
+          {
+            score: 0.8,
+            title: 'bar',
+            state: {
+              preferredSeriesType: 'bar_stacked',
+              legend: {
+                isVisible: true,
+                position: 'right',
+              },
+            },
+            previewIcon: 'empty',
+            visualizationId: 'lnsXY',
+          },
+          {
+            score: 0.8,
+            title: 'Test2',
+            state: {},
+            previewIcon: 'empty',
+          },
+          {
+            score: 0.8,
+            title: 'Test2',
+            state: {},
+            previewIcon: 'empty',
+            incomplete: true,
+          },
+        ],
+      },
+    };
+    datasourceMap.textBased.getDatasourceSuggestionsForVisualizeField.mockReturnValue([
+      generateSuggestion(),
+    ]);
+    datasourceMap.textBased.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      generateSuggestion(),
+    ]);
+    const context = {
+      dataViewSpec: {
+        id: 'index1',
+        title: 'index1',
+        name: 'DataView',
+      },
+      fieldName: '',
+      textBasedColumns: textBasedQueryColumns,
+      query: {
+        esql: 'FROM "index1" | keep field1, field2',
+      },
+    };
+    const suggestions = suggestionsApi({
+      context,
+      dataView,
+      datasourceMap,
+      visualizationMap,
+      preferredChartType: ChartType.XY,
+      preferredVisAttributes: {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+            legend: {
+              isVisible: false,
+              position: 'left',
+            },
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'],
+    });
+    expect(suggestions?.length).toEqual(1);
+    expect(suggestions?.[0]).toMatchInlineSnapshot(`
+      Object {
+        "changeType": "unchanged",
+        "columns": 0,
+        "datasourceId": "textBased",
+        "datasourceState": Object {
+          "layers": Object {},
+        },
+        "keptLayerIds": Array [],
+        "previewIcon": "empty",
+        "score": 0.8,
+        "title": undefined,
+        "visualizationId": "lnsXY",
+        "visualizationState": Object {
+          "legend": Object {
+            "isVisible": false,
+            "position": "left",
+          },
+          "preferredSeriesType": "bar_stacked",
+        },
+      }
+    `);
   });
 
   test('filters out the suggestion if exists on excludedVisualizations', async () => {

x-pack/plugins/lens/public/lens_suggestions_api/readme.md

@@ -0,0 +1,77 @@
+# Lens Suggestions API
+
+This document provides an overview of the Lens Suggestions API. It is used mostly for suggesting ES|QL charts based on an ES|QL query. It is used by the observability assistant, Discover and Dashboards ES|QL charts.
+
+## Overview
+
+The Lens Suggestions API is designed to provide suggestions for visualizations based on a given ES|QL query. It helps users to quickly find the most relevant visualizations for their data.
+
+## Getting Started
+
+To use the Lens Suggestions API, you need to import it from the Lens plugin:
+
+```typescript
+import useAsync from 'react-use/lib/useAsync';
+
+const lensHelpersAsync = useAsync(() => {
+    return lensService?.stateHelperApi() ?? Promise.resolve(null);
+  }, [lensService]);
+
+  if (lensHelpersAsync.value) {
+    const suggestionsApi = lensHelpersAsync.value.suggestions;
+  }
+```
+
+## The api
+
+The api returns an array of suggestions.
+
+#### Parameters
+
+  dataView: DataView;
+  visualizationMap?: VisualizationMap;
+  datasourceMap?: DatasourceMap;
+  excludedVisualizations?: string[];
+  preferredChartType?: ChartType;
+  preferredVisAttributes?: TypedLensByValueInput['attributes'];
+
+- `context`: The context as descibed by the VisualizeFieldContext. 
+- `dataView`: The dataView, can be an adhoc one too. For ES|QL you can create a dataview like this
+
+```typescript
+const indexName = (await getIndexForESQLQuery({ dataViews })) ?? '*';
+const dataView = await getESQLAdHocDataview(`from ${indexName}`, dataViews);
+```
+Optional parameters:
+- `preferredChartType`: Use this if you want the suggestions api to prioritize a specific suggestion type.
+- `preferredVisAttributes`: Use this with the preferredChartType if you want to prioritize a specific suggestion type with a non-default visualization state.
+
+#### Returns
+
+An array of suggestion objects
+
+## Example Usage
+
+```typescript
+const abc = new AbortController();
+
+const columns = await getESQLQueryColumns({
+  esqlQuery,
+  search: dataService.search.search,
+  signal: abc.signal,
+  timeRange: dataService.query.timefilter.timefilter.getAbsoluteTime(),
+});
+
+const context = {
+  dataViewSpec: dataView?.toSpec(false),
+  fieldName: '',
+  textBasedColumns: columns,
+  query: { esql: esqlQuery },
+};
+
+const chartSuggestions = lensHelpersAsync.value.suggestions(context, dataView);
+
+suggestions.forEach(suggestion => {
+  console.log(`Suggestion: ${suggestion.title}, Score: ${suggestion.score}`);
+});
+```

x-pack/plugins/lens/public/plugin.ts

@@ -62,6 +62,7 @@ import {
   ContentManagementPublicStart,
 } from '@kbn/content-management-plugin/public';
 import { i18n } from '@kbn/i18n';
+import type { ChartType } from '@kbn/visualization-utils';
 import type { ServerlessPluginStart } from '@kbn/serverless/public';
 import { LicensingPluginStart } from '@kbn/licensing-plugin/public';
 import type { EditorFrameService as EditorFrameServiceType } from './editor_frame_service';
@@ -137,7 +138,7 @@ import {
 } from '../common/content_management';
 import type { EditLensConfigurationProps } from './app_plugin/shared/edit_on_the_fly/get_edit_lens_configuration';
 import { savedObjectToEmbeddableAttributes } from './lens_attribute_service';
-import { ChartType } from './lens_suggestions_api';
+import type { TypedLensByValueInput } from './embeddable/embeddable_component';
 
 export type { SaveProps } from './app_plugin';
 
@@ -281,7 +282,8 @@ export type LensSuggestionsApi = (
   context: VisualizeFieldContext | VisualizeEditorContext,
   dataViews: DataView,
   excludedVisualizations?: string[],
-  preferredChartType?: ChartType
+  preferredChartType?: ChartType,
+  preferredVisAttributes?: TypedLensByValueInput['attributes']
 ) => Suggestion[] | undefined;
 
 export class LensPlugin {
@@ -713,7 +715,13 @@ export class LensPlugin {
         return {
           formula: createFormulaPublicApi(),
           chartInfo: createChartInfoApi(startDependencies.dataViews, this.editorFrameService),
-          suggestions: (context, dataView, excludedVisualizations, preferredChartType) => {
+          suggestions: (
+            context,
+            dataView,
+            excludedVisualizations,
+            preferredChartType,
+            preferredVisAttributes
+          ) => {
             return suggestionsApi({
               datasourceMap,
               visualizationMap,
@@ -721,6 +729,7 @@ export class LensPlugin {
               dataView,
               excludedVisualizations,
               preferredChartType,
+              preferredVisAttributes,
             });
           },
         };

x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx

@@ -37,7 +37,7 @@ import {
   VisualizeESQLUserIntention,
 } from '@kbn/observability-ai-assistant-plugin/public';
 import type { UiActionsStart } from '@kbn/ui-actions-plugin/public';
-import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import { getLensAttributesFromSuggestion, ChartType } from '@kbn/visualization-utils';
 import React, { useCallback, useContext, useEffect, useMemo, useState } from 'react';
 import ReactDOM from 'react-dom';
 import useAsync from 'react-use/lib/useAsync';
@@ -48,19 +48,6 @@ import type {
 } from '../../common/functions/visualize_esql';
 import { ObservabilityAIAssistantAppPluginStartDependencies } from '../types';
 
-enum ChartType {
-  XY = 'XY',
-  Bar = 'Bar',
-  Line = 'Line',
-  Area = 'Area',
-  Donut = 'Donut',
-  Heatmap = 'Heat map',
-  Treemap = 'Treemap',
-  Tagcloud = 'Tag cloud',
-  Waffle = 'Waffle',
-  Table = 'Table',
-}
-
 interface VisualizeESQLProps {
   /** Lens start contract, get the ES|QL charts suggestions api */
   lens: LensPublicStart;