[Lens] Fix error in query from generated suggestion

x-pack/legacy/plugins/lens/public/indexpattern_datasource/indexpattern_suggestions.test.tsx

@@ -827,7 +827,7 @@ describe('IndexPattern Data Source suggestions', () => {
                   columnOrder: ['cola', 'id1'],
                   columns: {
                     cola: initialState.layers.currentLayer.columns.cola,
-                    id1: expect.objectContaining({
+                    colb: expect.objectContaining({
                       operationType: 'avg',
                       sourceField: 'memory',
                     }),
@@ -839,6 +839,16 @@ describe('IndexPattern Data Source suggestions', () => {
         );
       });
 
+      it('updates references when replacing one metric with another', () => {
+        const initialState = stateWithNonEmptyTables();
+        const suggestions = getDatasourceSuggestionsForField(initialState, '1', {
+          name: 'memory',
+          type: 'number',
+          aggregatable: true,
+          searchable: true,
+        });
+      });
+
       it('adds a metric column on a number field if no other metrics set', () => {
         const initialState = stateWithNonEmptyTables();
         const modifiedState: IndexPatternPrivateState = {

x-pack/legacy/plugins/lens/public/indexpattern_datasource/indexpattern_suggestions.ts

@@ -15,8 +15,8 @@ import {
   operationDefinitionMap,
   IndexPatternColumn,
 } from './operations';
-import { hasField } from './utils';
 import { operationDefinitions } from './operations/definitions';
+import { hasField } from './utils';
 import {
   IndexPattern,
   IndexPatternPrivateState,
@@ -196,30 +196,29 @@ function addFieldAsMetricOperation(
     suggestedPriority: undefined,
     field,
   });
-  const newColumnId = generateId();
+  const addedColumnId = generateId();
 
-  const [, metrics] = separateBucketColumns(layer);
+  const [buckets, metrics] = separateBucketColumns(layer);
 
   // Add metrics if there are 0 or > 1 metric
   if (metrics.length !== 1) {
     return {
       indexPatternId: indexPattern.id,
       columns: {
         ...layer.columns,
-        [newColumnId]: newColumn,
+        [addedColumnId]: newColumn,
       },
-      columnOrder: [...layer.columnOrder, newColumnId],
+      columnOrder: [...layer.columnOrder, addedColumnId],
     };
   }
 
-  // If only one metric, replace instead of add
-  const newColumns = { ...layer.columns, [newColumnId]: newColumn };
-  delete newColumns[metrics[0]];
+  // Replacing old column with new column, keeping the old ID
+  const newColumns = { ...layer.columns, [metrics[0]]: newColumn };
 
   return {
     indexPatternId: indexPattern.id,
     columns: newColumns,
-    columnOrder: [...layer.columnOrder.filter(c => c !== metrics[0]), newColumnId],
+    columnOrder: layer.columnOrder, // Order is kept by replacing
   };
 }
 

x-pack/legacy/plugins/lens/public/indexpattern_datasource/operations/definitions/terms.test.tsx

@@ -274,7 +274,7 @@ describe('terms', () => {
       expect(updatedColumn).toBe(initialColumn);
     });
 
-    it('should switch to alphabetical ordering if the order column is removed', () => {
+    it('should switch to alphabetical ordering if there are no columns to order by', () => {
       const termsColumn = termsOperation.onOtherColumnChanged!(
         {
           label: 'Top value of category',

x-pack/legacy/plugins/lens/public/indexpattern_datasource/operations/definitions/terms.tsx

@@ -135,7 +135,7 @@ export const termsOperation: OperationDefinition<TermsIndexPatternColumn> = {
     }
     return currentColumn;
   },
-  paramEditor: ({ state, setState, currentColumn, columnId: currentColumnId, layerId }) => {
+  paramEditor: ({ state, setState, currentColumn, layerId }) => {
     const SEPARATOR = '$$$';
     function toValue(orderBy: TermsIndexPatternColumn['params']['orderBy']) {
       if (orderBy.type === 'alphabetical') {