[Lens] Create mathColumn function to improve performance (#101908)

* [Lens] Create mathColumn function to improve performance

* Fix empty formula case

* Fix tinymath memoization

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

docs/canvas/canvas-function-reference.asciidoc

@@ -71,7 +71,7 @@ Alias: `condition`
 [[alterColumn_fn]]
 === `alterColumn`
 
-Converts between core types, including `string`, `number`, `null`, `boolean`, and `date`, and renames columns. See also <<mapColumn_fn>> and <<staticColumn_fn>>.
+Converts between core types, including `string`, `number`, `null`, `boolean`, and `date`, and renames columns. See also <<mapColumn_fn>>, <<mathColumn_fn>>, and <<staticColumn_fn>>.
 
 *Expression syntax*
 [source,js]
@@ -1717,11 +1717,16 @@ Adds a column calculated as the result of other columns. Changes are made only w
 |===
 |Argument |Type |Description
 
+|`id`
+
+|`string`, `null`
+|An optional id of the resulting column. When no id is provided, the id will be looked up from the existing column by the provided name argument. If no column with this name exists yet, a new column with this name and an identical id will be added to the table.
+
 |_Unnamed_ ***
 
 Aliases: `column`, `name`
 |`string`
-|The name of the resulting column.
+|The name of the resulting column. Names are not required to be unique.
 
 |`expression` ***
 
@@ -1729,11 +1734,6 @@ Aliases: `exp`, `fn`, `function`
 |`boolean`, `number`, `string`, `null`
 |A Canvas expression that is passed to each row as a single row `datatable`.
 
-|`id`
-
-|`string`, `null`
-|An optional id of the resulting column. When not specified or `null` the name argument is used as id.
-
 |`copyMetaFrom`
 
 |`string`, `null`
@@ -1808,6 +1808,47 @@ Default: `"throw"`
 *Returns:* `number` | `boolean` | `null`
 
 
+[float]
+[[mathColumn_fn]]
+=== `mathColumn`
+
+Adds a column by evaluating `TinyMath` on each row. This function is optimized for math, so it performs better than the <<mapColumn_fn>> with a <<math_fn>>.
+*Accepts:* `datatable`
+
+[cols="3*^<"]
+|===
+|Argument |Type |Description
+
+|id ***
+|`string`
+|id of the resulting column. Must be unique.
+
+|name ***
+|`string`
+|The name of the resulting column. Names are not required to be unique.
+
+|_Unnamed_
+
+Alias: `expression`
+|`string`
+|A `TinyMath` expression evaluated on each row. See https://www.elastic.co/guide/en/kibana/current/canvas-tinymath-functions.html.
+
+|`onError`
+
+|`string`
+|In case the `TinyMath` evaluation fails or returns NaN, the return value is specified by onError. For example, `"null"`, `"zero"`, `"false"`, `"throw"`. When `"throw"`, it will throw an exception, terminating expression execution.
+
+Default: `"throw"`
+
+|`copyMetaFrom`
+
+|`string`, `null`
+|If set, the meta object from the specified column id is copied over to the specified target column. Throws an exception if the column doesn't exist
+|===
+
+*Returns:* `datatable`
+
+
 [float]
 [[metric_fn]]
 === `metric`
@@ -2581,7 +2622,7 @@ Default: `false`
 [[staticColumn_fn]]
 === `staticColumn`
 
-Adds a column with the same static value in every row. See also <<alterColumn_fn>> and <<mapColumn_fn>>.
+Adds a column with the same static value in every row. See also <<alterColumn_fn>>, <<mapColumn_fn>>, and <<mathColumn_fn>>.
 
 *Accepts:* `datatable`
 

packages/kbn-tinymath/src/index.js

@@ -7,12 +7,11 @@
  */
 
 const { get } = require('lodash');
+const memoizeOne = require('memoize-one');
 // eslint-disable-next-line import/no-unresolved
 const { parse: parseFn } = require('../grammar');
 const { functions: includedFunctions } = require('./functions');
 
-module.exports = { parse, evaluate, interpret };
-
 function parse(input, options) {
   if (input == null) {
     throw new Error('Missing expression');
@@ -29,9 +28,11 @@ function parse(input, options) {
   }
 }
 
+const memoizedParse = memoizeOne(parse);
+
 function evaluate(expression, scope = {}, injectedFunctions = {}) {
   scope = scope || {};
-  return interpret(parse(expression), scope, injectedFunctions);
+  return interpret(memoizedParse(expression), scope, injectedFunctions);
 }
 
 function interpret(node, scope, injectedFunctions) {
@@ -79,3 +80,5 @@ function isOperable(args) {
     return typeof arg === 'number' && !isNaN(arg);
   });
 }
+
+module.exports = { parse: memoizedParse, evaluate, interpret };

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

@@ -18,3 +18,4 @@ export * from './moving_average';
 export * from './ui_setting';
 export { mapColumn, MapColumnArguments } from './map_column';
 export { math, MathArguments, MathInput } from './math';
+export { mathColumn, MathColumnArguments } from './math_column';

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

@@ -0,0 +1,111 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+import { i18n } from '@kbn/i18n';
+import { ExpressionFunctionDefinition } from '../types';
+import { math, MathArguments } from './math';
+import { Datatable, DatatableColumn, getType } from '../../expression_types';
+
+export type MathColumnArguments = MathArguments & {
+  id: string;
+  name?: string;
+  copyMetaFrom?: string | null;
+};
+
+export const mathColumn: ExpressionFunctionDefinition<
+  'mathColumn',
+  Datatable,
+  MathColumnArguments,
+  Datatable
+> = {
+  name: 'mathColumn',
+  type: 'datatable',
+  inputTypes: ['datatable'],
+  help: i18n.translate('expressions.functions.mathColumnHelpText', {
+    defaultMessage:
+      'Adds a column calculated as the result of other columns. ' +
+      'Changes are made only when you provide arguments.' +
+      'See also {alterColumnFn} and {staticColumnFn}.',
+    values: {
+      alterColumnFn: '`alterColumn`',
+      staticColumnFn: '`staticColumn`',
+    },
+  }),
+  args: {
+    ...math.args,
+    id: {
+      types: ['string'],
+      help: i18n.translate('expressions.functions.mathColumn.args.idHelpText', {
+        defaultMessage: 'id of the resulting column. Must be unique.',
+      }),
+      required: true,
+    },
+    name: {
+      types: ['string'],
+      aliases: ['_', 'column'],
+      help: i18n.translate('expressions.functions.mathColumn.args.nameHelpText', {
+        defaultMessage: 'The name of the resulting column. Names are not required to be unique.',
+      }),
+      required: true,
+    },
+    copyMetaFrom: {
+      types: ['string', 'null'],
+      help: i18n.translate('expressions.functions.mathColumn.args.copyMetaFromHelpText', {
+        defaultMessage:
+          "If set, the meta object from the specified column id is copied over to the specified target column. If the column doesn't exist it silently fails.",
+      }),
+      required: false,
+      default: null,
+    },
+  },
+  fn: (input, args, context) => {
+    const columns = [...input.columns];
+    const existingColumnIndex = columns.findIndex(({ id }) => {
+      return id === args.id;
+    });
+    if (existingColumnIndex > -1) {
+      throw new Error('ID must be unique');
+    }
+
+    const newRows = input.rows.map((row) => {
+      return {
+        ...row,
+        [args.id]: math.fn(
+          {
+            type: 'datatable',
+            columns: input.columns,
+            rows: [row],
+          },
+          {
+            expression: args.expression,
+            onError: args.onError,
+          },
+          context
+        ),
+      };
+    });
+    const type = newRows.length ? getType(newRows[0][args.id]) : 'null';
+    const newColumn: DatatableColumn = {
+      id: args.id,
+      name: args.name ?? args.id,
+      meta: { type, params: { id: type } },
+    };
+    if (args.copyMetaFrom) {
+      const metaSourceFrom = columns.find(({ id }) => id === args.copyMetaFrom);
+      newColumn.meta = { ...newColumn.meta, ...(metaSourceFrom?.meta || {}) };
+    }
+
+    columns.push(newColumn);
+
+    return {
+      type: 'datatable',
+      columns,
+      rows: newRows,
+    } as Datatable;
+  },
+};

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

@@ -0,0 +1,74 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+import { mathColumn } from '../math_column';
+import { functionWrapper, testTable } from './utils';
+
+describe('mathColumn', () => {
+  const fn = functionWrapper(mathColumn);
+
+  it('throws if the id is used', () => {
+    expect(() => fn(testTable, { id: 'price', name: 'price', expression: 'price * 2' })).toThrow(
+      `ID must be unique`
+    );
+  });
+
+  it('applies math to each row by id', () => {
+    const result = fn(testTable, { id: 'output', name: 'output', expression: 'quantity * price' });
+    expect(result.columns).toEqual([
+      ...testTable.columns,
+      { id: 'output', name: 'output', meta: { params: { id: 'number' }, type: 'number' } },
+    ]);
+    expect(result.rows[0]).toEqual({
+      in_stock: true,
+      name: 'product1',
+      output: 60500,
+      price: 605,
+      quantity: 100,
+      time: 1517842800950,
+    });
+  });
+
+  it('handles onError', () => {
+    const args = {
+      id: 'output',
+      name: 'output',
+      expression: 'quantity / 0',
+    };
+    expect(() => fn(testTable, args)).toThrowError(`Cannot divide by 0`);
+    expect(() => fn(testTable, { ...args, onError: 'throw' })).toThrow();
+    expect(fn(testTable, { ...args, onError: 'zero' }).rows[0].output).toEqual(0);
+    expect(fn(testTable, { ...args, onError: 'false' }).rows[0].output).toEqual(false);
+    expect(fn(testTable, { ...args, onError: 'null' }).rows[0].output).toEqual(null);
+  });
+
+  it('should copy over the meta information from the specified column', async () => {
+    const result = await fn(
+      {
+        ...testTable,
+        columns: [
+          ...testTable.columns,
+          {
+            id: 'myId',
+            name: 'myName',
+            meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+          },
+        ],
+        rows: testTable.rows.map((row) => ({ ...row, myId: Date.now() })),
+      },
+      { id: 'output', name: 'name', copyMetaFrom: 'myId', expression: 'price + 2' }
+    );
+
+    expect(result.type).toBe('datatable');
+    expect(result.columns[result.columns.length - 1]).toEqual({
+      id: 'output',
+      name: 'name',
+      meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+    });
+  });
+});

src/plugins/expressions/common/service/expressions_services.ts

@@ -31,6 +31,7 @@ import {
   mapColumn,
   overallMetric,
   math,
+  mathColumn,
 } from '../expression_functions';
 
 /**
@@ -344,6 +345,7 @@ export class ExpressionsService implements PersistableStateService<ExpressionAst
       overallMetric,
       mapColumn,
       math,
+      mathColumn,
     ]) {
       this.registerFunction(fn);
     }

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

@@ -100,28 +100,11 @@ export const formulaOperation: OperationDefinition<
     return [
       {
         type: 'function',
-        function: 'mapColumn',
+        function: currentColumn.references.length ? 'mathColumn' : 'mapColumn',
         arguments: {
           id: [columnId],
           name: [label || defaultLabel],
-          exp: [
-            {
-              type: 'expression',
-              chain: currentColumn.references.length
-                ? [
-                    {
-                      type: 'function',
-                      function: 'math',
-                      arguments: {
-                        expression: [
-                          currentColumn.references.length ? `"${currentColumn.references[0]}"` : ``,
-                        ],
-                      },
-                    },
-                  ]
-                : [],
-            },
-          ],
+          expression: [currentColumn.references.length ? `"${currentColumn.references[0]}"` : ''],
         },
       },
     ];

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

@@ -45,25 +45,12 @@ export const mathOperation: OperationDefinition<MathIndexPatternColumn, 'managed
     return [
       {
         type: 'function',
-        function: 'mapColumn',
+        function: 'mathColumn',
         arguments: {
           id: [columnId],
           name: [columnId],
-          exp: [
-            {
-              type: 'expression',
-              chain: [
-                {
-                  type: 'function',
-                  function: 'math',
-                  arguments: {
-                    expression: [astToString(column.params.tinymathAst)],
-                    onError: ['null'],
-                  },
-                },
-              ],
-            },
-          ],
+          expression: [astToString(column.params.tinymathAst)],
+          onError: ['null'],
         },
       },
     ];