Keep preferred type in transformational commands too

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

@@ -23,14 +23,11 @@ import type {
   Suggestion,
   TermsIndexPatternColumn,
   TypedLensByValueInput,
-  PieVisualizationState,
-  XYState,
 } from '@kbn/lens-plugin/public';
 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 type { ChartType } 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';
@@ -50,6 +47,7 @@ import { computeInterval } from '../utils/compute_interval';
 import { fieldSupportsBreakdown } from '../utils/field_supports_breakdown';
 import { shouldDisplayHistogram } from '../layout/helpers';
 import { enrichLensAttributesWithTablesData } from '../utils/lens_vis_from_table';
+import { getPreferredChartType } from '../utils/get_preferred_chart_type';
 
 const UNIFIED_HISTOGRAM_LAYER_ID = 'unifiedHistogram';
 
@@ -150,7 +148,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,
@@ -516,25 +517,12 @@ export class LensVisService {
       if (breakdownColumn) {
         context.textBasedColumns.push(breakdownColumn);
       }
-      let preferredChartType = preferredVisAttributes
-        ? preferredVisAttributes?.visualizationType
+      const preferredChartType = preferredVisAttributes
+        ? getPreferredChartType(preferredVisAttributes)
         : undefined;
 
-      if (preferredChartType === 'lnsXY') {
-        preferredChartType = (preferredVisAttributes?.state?.visualization as XYState)
-          ?.preferredSeriesType;
-      }
-      if (preferredChartType === 'lnsPie') {
-        preferredChartType = (preferredVisAttributes?.state?.visualization as PieVisualizationState)
-          ?.shape;
-      }
       const suggestions =
-        this.lensSuggestionsApi(
-          context,
-          dataView,
-          ['lnsDatatable'],
-          preferredChartType as ChartType
-        ) ?? [];
+        this.lensSuggestionsApi(context, dataView, ['lnsDatatable'], preferredChartType) ?? [];
       if (suggestions.length) {
         const suggestion = suggestions[0];
         const suggestionVisualizationState = Object.assign({}, suggestion?.visualizationState);
@@ -598,17 +586,27 @@ 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
+      ? getPreferredChartType(preferredVisAttributes)
+      : undefined;
+
     const context = {
       dataViewSpec: dataView?.toSpec(),
       fieldName: '',
       textBasedColumns: columns,
       query: query && isOfAggregateQueryType(query) ? query : undefined,
     };
     const allSuggestions = isPlainRecord
-      ? this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? []
+      ? this.lensSuggestionsApi(context, dataView, ['lnsDatatable'], preferredChartType) ?? []
       : [];
 
     return allSuggestions;

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

@@ -0,0 +1,24 @@
+/*
+ * 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 { PieVisualizationState, XYState } from '@kbn/lens-plugin/public';
+import type { ChartType } from '@kbn/visualization-utils';
+import type { UnifiedHistogramVisContext } from '../types';
+
+export const getPreferredChartType = (visAttributes: UnifiedHistogramVisContext['attributes']) => {
+  let preferredChartType = visAttributes ? visAttributes?.visualizationType : undefined;
+
+  if (preferredChartType === 'lnsXY') {
+    preferredChartType = (visAttributes?.state?.visualization as XYState)?.preferredSeriesType;
+  }
+  if (preferredChartType === 'lnsPie') {
+    preferredChartType = (visAttributes?.state?.visualization as PieVisualizationState)?.shape;
+  }
+
+  return preferredChartType as ChartType;
+};