[Lens] Allow date functions in formula (#143632)

* allow date functions in formula

* fix tests

* fix test

* fix review comments

* fix test

Co-authored-by: Marco Liberati <dej611@users.noreply.github.com>

src/plugins/expressions/common/expression_functions/specs/math_column.ts

@@ -14,6 +14,7 @@ import { Datatable, DatatableColumn, DatatableColumnType, getType } from '../../
 export type MathColumnArguments = MathArguments & {
   id: string;
   name?: string;
+  castColumns?: string[];
   copyMetaFrom?: string | null;
 };
 
@@ -52,6 +53,14 @@ export const mathColumn: ExpressionFunctionDefinition<
       }),
       required: true,
     },
+    castColumns: {
+      types: ['string'],
+      multi: true,
+      help: i18n.translate('expressions.functions.mathColumn.args.castColumnsHelpText', {
+        defaultMessage: 'The ids of columns to cast to numbers before applying the formula',
+      }),
+      required: false,
+    },
     copyMetaFrom: {
       types: ['string', 'null'],
       help: i18n.translate('expressions.functions.mathColumn.args.copyMetaFromHelpText', {
@@ -77,11 +86,31 @@ export const mathColumn: ExpressionFunctionDefinition<
 
     const newRows = await Promise.all(
       input.rows.map(async (row) => {
+        let preparedRow = row;
+        if (args.castColumns) {
+          preparedRow = { ...row };
+          args.castColumns.forEach((columnId) => {
+            switch (typeof row[columnId]) {
+              case 'string':
+                const parsedAsDate = Number(new Date(preparedRow[columnId]));
+                if (!isNaN(parsedAsDate)) {
+                  preparedRow[columnId] = parsedAsDate;
+                  return;
+                } else {
+                  preparedRow[columnId] = Number(preparedRow[columnId]);
+                  return;
+                }
+              case 'boolean':
+                preparedRow[columnId] = Number(preparedRow[columnId]);
+                return;
+            }
+          });
+        }
         const result = await math.fn(
           {
             ...input,
             columns: input.columns,
-            rows: [row],
+            rows: [preparedRow],
           },
           {
             expression: args.expression,

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/editor/math_completion.test.ts

@@ -353,7 +353,7 @@ describe('math completion', () => {
       expect(results.list).toEqual(['bytes', 'memory']);
     });
 
-    it('should autocomplete only operations that provide numeric output', async () => {
+    it('should autocomplete only operations that provide numeric or date output', async () => {
       const results = await suggest({
         expression: 'last_value()',
         zeroIndexedOffset: 11,
@@ -366,7 +366,7 @@ describe('math completion', () => {
         unifiedSearch: unifiedSearchPluginMock.createStartContract(),
         dataViews: dataViewPluginMocks.createStartContract(),
       });
-      expect(results.list).toEqual(['bytes', 'memory']);
+      expect(results.list).toEqual(['bytes', 'memory', 'timestamp', 'start_date']);
     });
   });
 

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/editor/math_completion.ts

@@ -306,12 +306,14 @@ function getArgumentSuggestions(
       operationDefinitionMap
     );
     // TODO: This only allow numeric functions, will reject last_value(string) for example.
-    const validOperation = available.find(
+    const validOperation = available.filter(
       ({ operationMetaData }) =>
-        operationMetaData.dataType === 'number' && !operationMetaData.isBucketed
+        (operationMetaData.dataType === 'number' || operationMetaData.dataType === 'date') &&
+        !operationMetaData.isBucketed
     );
-    if (validOperation) {
+    if (validOperation.length) {
-      const fields = validOperation.operations
+      const fields = validOperation
+        .flatMap((op) => op.operations)
         .filter((op) => op.operationType === operation.type)
         .map((op) => ('field' in op ? op.field : undefined))
         .filter(nonNullable);

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/formula.test.tsx

@@ -157,7 +157,7 @@ describe('formula', () => {
         formulaOperation.buildColumn({
           previousColumn: {
             ...layer.columns.col1,
-            dataType: 'date',
+            dataType: 'boolean',
             filter: { language: 'kuery', query: 'ABC: DEF' },
           },
           layer,

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/formula.tsx

@@ -139,6 +139,11 @@ export const formulaOperation: OperationDefinition<FormulaIndexPatternColumn, 'm
           arguments: {
             id: [columnId],
             name: [label || defaultLabel],
+            ...(currentColumn.references.length
+              ? {
+                  castColumns: [currentColumn.references[0]],
+                }
+              : {}),
             expression: [currentColumn.references.length ? `"${currentColumn.references[0]}"` : ''],
           },
         },

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/generate.ts

@@ -54,7 +54,11 @@ export function generateFormula(
       previousFormula += `${previousColumn.operationType}(${metric.operationType}(${fieldName})`;
     }
   } else {
-    if (previousColumn && 'sourceField' in previousColumn && previousColumn.dataType === 'number') {
+    if (
+      previousColumn &&
+      'sourceField' in previousColumn &&
+      (previousColumn.dataType === 'number' || previousColumn.dataType === 'date')
+    ) {
       previousFormula += `${previousColumn.operationType}(${getSafeFieldName(previousColumn)}`;
     } else {
       // couldn't find formula function to call, exit early because adding args is going to fail anyway

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/math.test.ts

@@ -163,6 +163,7 @@ describe('math operation', () => {
           arguments: {
             id: ['myColumnId'],
             name: ['Math'],
+            castColumns: [],
             expression: [
               '(((((((((((((((("columnX0" + "columnX1") + "columnX2") + "columnX3") + "columnX4") + "columnX5") + "columnX6") + "columnX7") + "columnX8") + "columnX9") + "columnX10") + "columnX11") + "columnX12") + "columnX13") + "columnX14") + "columnX15") + "columnX16")',
             ],
@@ -243,6 +244,7 @@ describe('math operation', () => {
           arguments: {
             id: ['myColumnId'],
             name: ['Math'],
+            castColumns: [],
             expression: [
               `("columnX0" + (("columnX1" - "columnX2") / ("columnX3" - ("columnX4" * "columnX5"))))`,
             ],
@@ -298,6 +300,7 @@ describe('math operation', () => {
           arguments: {
             id: ['myColumnId'],
             name: ['Math'],
+            castColumns: [],
             expression: [`max(min("columnX0","columnX1"),abs("columnX2"))`],
             onError: ['null'],
           },
@@ -342,6 +345,7 @@ describe('math operation', () => {
           arguments: {
             id: ['myColumnId'],
             name: ['Math'],
+            castColumns: [],
             expression: [`(5 + (3 / 8))`],
             onError: ['null'],
           },
@@ -425,6 +429,7 @@ describe('math operation', () => {
           arguments: {
             id: ['myColumnId'],
             name: ['Math'],
+            castColumns: [],
             expression: [
               'ifelse(("columnX0" == 0),ifelse(("columnX1" < 0),ifelse(("columnX2" <= 0),"columnX3","columnX4"),"columnX5"),ifelse(("columnX6" > 0),ifelse(("columnX7" >= 0),"columnX8","columnX9"),"columnX10"))',
             ],

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/formula/math.tsx

@@ -46,6 +46,8 @@ export const mathOperation: OperationDefinition<MathIndexPatternColumn, 'managed
           name: [column.label],
           expression: [astToString(column.params.tinymathAst)],
           onError: ['null'],
+          // cast everything into number
+          castColumns: column.references,
         },
       },
     ];

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/last_value.test.tsx

@@ -319,7 +319,7 @@ describe('last_value', () => {
       ).toEqual({
         dataType: 'ip',
         isBucketed: false,
-        scale: 'ratio',
+        scale: 'ordinal',
       });
     });
 

x-pack/plugins/lens/public/datasources/form_based/operations/definitions/last_value.tsx

@@ -124,6 +124,16 @@ function getExistsFilter(field: string) {
   };
 }
 
+function getScale(type: string) {
+  return type === 'string' ||
+    type === 'ip' ||
+    type === 'ip_range' ||
+    type === 'date_range' ||
+    type === 'number_range'
+    ? 'ordinal'
+    : 'ratio';
+}
+
 export const lastValueOperation: OperationDefinition<
   LastValueIndexPatternColumn,
   'field',
@@ -155,7 +165,7 @@ export const lastValueOperation: OperationDefinition<
       label: ofName(field.displayName, oldColumn.timeShift, oldColumn.reducedTimeRange),
       sourceField: field.name,
       params: newParams,
-      scale: field.type === 'string' ? 'ordinal' : 'ratio',
+      scale: getScale(field.type),
       filter:
         oldColumn.filter && isEqual(oldColumn.filter, getExistsFilter(oldColumn.sourceField))
           ? getExistsFilter(field.name)
@@ -167,7 +177,7 @@ export const lastValueOperation: OperationDefinition<
       return {
         dataType: type as DataType,
         isBucketed: false,
-        scale: type === 'string' ? 'ordinal' : 'ratio',
+        scale: getScale(type),
       };
     }
   },
@@ -218,7 +228,7 @@ export const lastValueOperation: OperationDefinition<
       dataType: field.type as DataType,
       operationType: 'last_value',
       isBucketed: false,
-      scale: field.type === 'string' ? 'ordinal' : 'ratio',
+      scale: getScale(field.type),
       sourceField: field.name,
       filter: getFilter(previousColumn, columnParams) || getExistsFilter(field.name),
       timeShift: columnParams?.shift || previousColumn?.timeShift,