[Lens] optimize duplicate formula functions

src/plugins/data/public/index.ts

@@ -139,6 +139,8 @@ export type {
   ParsedInterval,
   // expressions
   ExecutionContextSearch,
+  ExpressionFunctionKql,
+  ExpressionFunctionLucene,
   ExpressionFunctionKibana,
   ExpressionFunctionKibanaContext,
   ExpressionValueSearchContext,

x-pack/plugins/lens/public/indexpattern_datasource/dedupe_aggs.test.ts

@@ -0,0 +1,108 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+  buildExpression,
+  ExpressionAstExpressionBuilder,
+  parseExpression,
+} from '@kbn/expressions-plugin/common';
+import { dedupeAggs } from './dedupe_aggs';
+import { operationDefinitionMap } from './operations';
+import { OriginalColumn } from './to_expression';
+
+describe('dedupeAggs', () => {
+  const buildMapsFromAggBuilders = (aggs: ExpressionAstExpressionBuilder[]) => {
+    const esAggsIdMap: Record<string, OriginalColumn[]> = {};
+    const aggsToIdsMap = new Map();
+    aggs.forEach((builder, i) => {
+      const esAggsId = `col-${i}-${i}`;
+      esAggsIdMap[esAggsId] = [{ id: `original-${i}` } as OriginalColumn];
+      aggsToIdsMap.set(builder, esAggsId);
+    });
+    return {
+      esAggsIdMap,
+      aggsToIdsMap,
+    };
+  };
+
+  it('removes duplicate aggregations', () => {
+    const aggs = [
+      'aggSum id="0" enabled=true schema="metric" field="bytes" emptyAsNull=false',
+      'aggSum id="1" enabled=true schema="metric" field="bytes" emptyAsNull=false',
+      'aggFilteredMetric id="2" enabled=true schema="metric" \n  customBucket={aggFilter id="2-filter" enabled=true schema="bucket" filter={kql q="hour_of_day: *"}} \n  customMetric={aggTopMetrics id="2-metric" enabled=true schema="metric" field="hour_of_day" size=1 sortOrder="desc" sortField="timestamp"}',
+      'aggFilteredMetric id="3" enabled=true schema="metric" \n  customBucket={aggFilter id="3-filter" enabled=true schema="bucket" filter={kql q="hour_of_day: *"}} \n  customMetric={aggTopMetrics id="3-metric" enabled=true schema="metric" field="hour_of_day" size=1 sortOrder="desc" sortField="timestamp"}',
+      'aggAvg id="4" enabled=true schema="metric" field="bytes"',
+      'aggAvg id="5" enabled=true schema="metric" field="bytes"',
+    ].map((expression) => buildExpression(parseExpression(expression)));
+
+    const { esAggsIdMap, aggsToIdsMap } = buildMapsFromAggBuilders(aggs);
+
+    // eslint-disable-next-line @typescript-eslint/naming-convention
+    const { sum, last_value, average } = operationDefinitionMap;
+
+    const operations = [sum, last_value, average];
+
+    operations.forEach((op) => expect(op.getGroupByKey).toBeDefined());
+
+    const { esAggsIdMap: newIdMap, aggs: newAggs } = dedupeAggs(
+      aggs,
+      esAggsIdMap,
+      aggsToIdsMap,
+      operations
+    );
+
+    expect(newAggs).toHaveLength(3);
+
+    expect(newIdMap).toMatchInlineSnapshot(`
+      Object {
+        "col-0-0": Array [
+          Object {
+            "id": "original-0",
+          },
+          Object {
+            "id": "original-1",
+          },
+        ],
+        "col-2-2": Array [
+          Object {
+            "id": "original-2",
+          },
+          Object {
+            "id": "original-3",
+          },
+        ],
+        "col-4-4": Array [
+          Object {
+            "id": "original-4",
+          },
+          Object {
+            "id": "original-5",
+          },
+        ],
+      }
+    `);
+  });
+
+  it('should update any terms order-by reference', () => {
+    const aggs = [
+      'aggTerms id="0" enabled=true schema="segment" field="clientip" orderBy="3" order="desc" size=5 includeIsRegex=false excludeIsRegex=false otherBucket=true otherBucketLabel="Other" missingBucket=false missingBucketLabel="(missing value)"',
+      'aggMedian id="1" enabled=true schema="metric" field="bytes"',
+      'aggMedian id="2" enabled=true schema="metric" field="bytes"',
+      'aggMedian id="3" enabled=true schema="metric" field="bytes"',
+    ].map((expression) => buildExpression(parseExpression(expression)));
+
+    const { esAggsIdMap, aggsToIdsMap } = buildMapsFromAggBuilders(aggs);
+
+    const { aggs: newAggs } = dedupeAggs(aggs, esAggsIdMap, aggsToIdsMap, [
+      operationDefinitionMap.median,
+    ]);
+
+    expect(newAggs).toHaveLength(2);
+
+    expect(newAggs[0].functions[0].getArgument('orderBy')?.[0]).toBe('1');
+  });
+});

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

@@ -0,0 +1,98 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { AggFunctionsMapping } from '@kbn/data-plugin/public';
+import {
+  ExpressionAstExpressionBuilder,
+  ExpressionAstFunctionBuilder,
+} from '@kbn/expressions-plugin/common';
+import { GenericOperationDefinition } from './operations';
+import { extractAggId, OriginalColumn } from './to_expression';
+
+function groupByKey<T>(items: T[], getKey: (item: T) => string | undefined): Record<string, T[]> {
+  const groups: Record<string, T[]> = {};
+
+  items.forEach((item) => {
+    const key = getKey(item);
+    if (key) {
+      if (!(key in groups)) {
+        groups[key] = [];
+      }
+      groups[key].push(item);
+    }
+  });
+
+  return groups;
+}
+
+/**
+ * Consolidates duplicate agg expression builders to increase performance
+ */
+export function dedupeAggs(
+  _aggs: ExpressionAstExpressionBuilder[],
+  _esAggsIdMap: Record<string, OriginalColumn[]>,
+  aggExpressionToEsAggsIdMap: Map<ExpressionAstExpressionBuilder, string>,
+  allOperations: GenericOperationDefinition[]
+): {
+  aggs: ExpressionAstExpressionBuilder[];
+  esAggsIdMap: Record<string, OriginalColumn[]>;
+} {
+  let aggs = [..._aggs];
+  const esAggsIdMap = { ..._esAggsIdMap };
+
+  const aggsByArgs = groupByKey<ExpressionAstExpressionBuilder>(aggs, (expressionBuilder) => {
+    for (const operation of allOperations) {
+      const groupKey = operation.getGroupByKey?.(expressionBuilder);
+      if (groupKey) {
+        return `${operation.type}-${groupKey}`;
+      }
+    }
+  });
+
+  const termsFuncs = aggs
+    .map((agg) => agg.functions[0])
+    .filter((func) => func.name === 'aggTerms') as Array<
+    ExpressionAstFunctionBuilder<AggFunctionsMapping['aggTerms']>
+  >;
+
+  // collapse each group into a single agg expression builder
+  Object.values(aggsByArgs).forEach((expressionBuilders) => {
+    if (expressionBuilders.length <= 1) {
+      // don't need to optimize if there aren't more than one
+      return;
+    }
+
+    const [firstExpressionBuilder, ...restExpressionBuilders] = expressionBuilders;
+
+    // throw away all but the first expression builder
+    aggs = aggs.filter((aggBuilder) => !restExpressionBuilders.includes(aggBuilder));
+
+    const firstEsAggsId = aggExpressionToEsAggsIdMap.get(firstExpressionBuilder);
+    if (firstEsAggsId === undefined) {
+      throw new Error('Could not find current column ID for expression builder');
+    }
+
+    restExpressionBuilders.forEach((expressionBuilder) => {
+      const currentEsAggsId = aggExpressionToEsAggsIdMap.get(expressionBuilder);
+      if (currentEsAggsId === undefined) {
+        throw new Error('Could not find current column ID for expression builder');
+      }
+
+      esAggsIdMap[firstEsAggsId].push(...esAggsIdMap[currentEsAggsId]);
+
+      delete esAggsIdMap[currentEsAggsId];
+
+      termsFuncs.forEach((func) => {
+        if (func.getArgument('orderBy')?.[0] === extractAggId(currentEsAggsId)) {
+          func.replaceArgument('orderBy', [extractAggId(firstEsAggsId)]);
+        }
+      });
+    });
+  });
+
+  return { aggs, esAggsIdMap };
+}