Fix field selector and start working on suggestions

x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/reference_editor.tsx

@@ -19,6 +19,7 @@ import {
   deleteColumn,
   RequiredReference,
   isOperationAllowedAsReference,
+  FieldBasedIndexPatternColumn,
 } from '../operations';
 import { FieldSelect } from './field_select';
 import { hasField } from '../utils';
@@ -53,6 +54,8 @@ export function ReferenceEditor(props: ReferenceEditorProps) {
   const selectedOperationDefinition = column && operationDefinitionMap[column.operationType];
 
   const incompleteInfo = layer.incompleteColumns ? layer.incompleteColumns[columnId] : undefined;
+  const incompleteOperation = incompleteInfo?.operationType;
+  const incompleteField = incompleteInfo?.sourceField ?? null;
 
   // Basically the operation support matrix, but different validation
   const operationSupportMatrix: OperationSupportMatrix & {
@@ -226,13 +229,17 @@ export function ReferenceEditor(props: ReferenceEditorProps) {
               currentIndexPattern={currentIndexPattern}
               existingFields={existingFields}
               operationSupportMatrix={operationSupportMatrix}
-              selectedColumnOperationType={column && column.operationType}
-              selectedColumnSourceField={
-                column && hasField(column) ? column.sourceField : undefined
+              selectedOperationType={
+                // Allows operation to be selected before creating a valid column
+                column ? column.operationType : incompleteOperation
               }
-              incompatibleSelectedOperationType={
-                (incompleteInfo?.operationType as OperationType) ?? null
+              selectedField={
+                // Allows field to be selected
+                incompleteField
+                  ? incompleteField
+                  : (column as FieldBasedIndexPatternColumn)?.sourceField
               }
+              incompleteOperation={incompleteOperation}
               onDeleteColumn={() => {
                 updateLayer(deleteColumn({ layer, columnId }));
               }}

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

@@ -1940,6 +1940,12 @@ describe('IndexPattern Data Source suggestions', () => {
       const suggestions = getDatasourceSuggestionsFromCurrentState(state);
       expect(suggestions).toEqual([]);
     });
+
+    describe('references', () => {
+      it('does not simplify a reference based operation to an invalid state', () => {
+        throw new Error('todo: write all the tests for references');
+      });
+    });
   });
 });
 

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

@@ -17,6 +17,7 @@ import {
   operationDefinitionMap,
   IndexPatternColumn,
   OperationType,
+  getExistingColumnGroups,
 } from './operations';
 import { hasField, hasInvalidFields } from './utils';
 import {
@@ -221,7 +222,7 @@ function getExistingLayerSuggestionsForField(
         );
       }
 
-      const [, metrics] = separateBucketColumns(layer);
+      const [, metrics, references] = getExistingColumnGroups(layer);
       if (metrics.length === 1) {
         const layerWithReplacedMetric = replaceColumn({
           layer,
@@ -377,7 +378,7 @@ export function getDatasourceSuggestionsFromCurrentState(
       .filter(([_id, layer]) => layer.columnOrder.length && layer.indexPatternId)
       .map(([layerId, layer]) => {
         const indexPattern = state.indexPatterns[layer.indexPatternId];
-        const [buckets, metrics] = separateBucketColumns(layer);
+        const [buckets, metrics, references] = getExistingColumnGroups(layer);
         const timeDimension = layer.columnOrder.find(
           (columnId) =>
             layer.columns[columnId].isBucketed && layer.columns[columnId].dataType === 'date'
@@ -570,7 +571,11 @@ function createSuggestionWithDefaultDateHistogram(
 function createSimplifiedTableSuggestions(state: IndexPatternPrivateState, layerId: string) {
   const layer = state.layers[layerId];
 
-  const [availableBucketedColumns, availableMetricColumns] = separateBucketColumns(layer);
+  const [
+    availableBucketedColumns,
+    availableMetricColumns,
+    availableReferenceColumns,
+  ] = getExistingColumnGroups(layer);
 
   return _.flatten(
     availableBucketedColumns.map((_col, index) => {
@@ -623,7 +628,3 @@ function getMetricSuggestionTitle(layer: IndexPatternLayer, onlyMetric: boolean)
       'Title of a suggested chart containing only a single numerical metric calculated over all available data',
   });
 }
-
-function separateBucketColumns(layer: IndexPatternLayer) {
-  return partition(layer.columnOrder, (columnId) => layer.columns[columnId].isBucketed);
-}

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/calculations/counter_rate.tsx

@@ -60,7 +60,8 @@ export const counterRateOperation: OperationDefinition<
     };
   },
   getDefaultLabel: (column, indexPattern, columns) => {
-    return ofName(columns[column.references[0]]?.label, column.timeScale);
+    const ref = columns[column.references[0]];
+    return ofName(ref && 'sourceField' in ref ? ref.sourceField : undefined, column.timeScale);
   },
   toExpression: (layer, columnId) => {
     return dateBasedOperationToExpression(layer, columnId, 'lens_counter_rate');
@@ -69,7 +70,7 @@ export const counterRateOperation: OperationDefinition<
     const metric = layer.columns[referenceIds[0]];
     const timeScale = previousColumn?.timeScale || DEFAULT_TIME_SCALE;
     return {
-      label: ofName(metric?.label, timeScale),
+      label: ofName(metric && 'sourceField' in metric ? metric.sourceField : undefined, timeScale),
       dataType: 'number',
       operationType: 'counter_rate',
       isBucketed: false,

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/calculations/derivative.tsx

@@ -59,15 +59,19 @@ export const derivativeOperation: OperationDefinition<
     };
   },
   getDefaultLabel: (column, indexPattern, columns) => {
-    return ofName(columns[column.references[0]]?.label, column.timeScale);
+    const ref = columns[column.references[0]];
+    return ofName(ref && 'sourceField' in ref ? ref.sourceField : undefined, column.timeScale);
   },
   toExpression: (layer, columnId) => {
     return dateBasedOperationToExpression(layer, columnId, 'derivative');
   },
   buildColumn: ({ referenceIds, previousColumn, layer }) => {
-    const metric = layer.columns[referenceIds[0]];
+    const ref = layer.columns[referenceIds[0]];
     return {
-      label: ofName(metric?.label, previousColumn?.timeScale),
+      label: ofName(
+        ref && 'sourceField' in ref ? ref.sourceField : undefined,
+        previousColumn?.timeScale
+      ),
       dataType: 'number',
       operationType: 'derivative',
       isBucketed: false,

x-pack/plugins/lens/public/indexpattern_datasource/operations/layer_helpers.ts

@@ -552,7 +552,7 @@ export function getColumnOrder(layer: IndexPatternLayer): string[] {
 }
 
 // Splits existing columnOrder into the three categories
-function getExistingColumnGroups(layer: IndexPatternLayer): [string[], string[], string[]] {
+export function getExistingColumnGroups(layer: IndexPatternLayer): [string[], string[], string[]] {
   const [direct, referenced] = partition(
     layer.columnOrder,
     (columnId) => layer.columns[columnId] && !('references' in layer.columns[columnId])