From d39084b24356f973f923c366230230bbeb8aa63e Mon Sep 17 00:00:00 2001
From: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
Date: Fri, 9 Sep 2022 08:15:07 -0600
Subject: [PATCH] apply single fn correctly (#140381) (#140393)

(cherry picked from commit 385dc10e3a790c63f9c32e67c76d52f5dfdc075a)

Co-authored-by: Joe Reuter <johannes.reuter@elastic.co>
---
 .../expressions/collapse/collapse_fn.test.ts  | 43 ++++++++++++++++---
 .../expressions/collapse/collapse_fn.ts       |  7 +--
 .../lens/common/expressions/collapse/index.ts |  2 +-
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.test.ts b/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.test.ts
index f878653430954..ed9f46f96b44c 100644
--- a/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.test.ts
+++ b/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.test.ts
@@ -33,12 +33,45 @@ describe('collapse_fn', () => {
           { val: 8, split: 'B' },
         ],
       },
-      { metric: ['val'], fn: 'sum' }
+      { metric: ['val'], fn: ['sum'] }
     );
 
     expect(result.rows).toEqual([{ val: 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 }]);
   });
 
+  it('can use a single function for multiple metrics', async () => {
+    const result = await runFn(
+      {
+        type: 'datatable',
+        columns: [
+          { id: 'val', name: 'val', meta: { type: 'number' } },
+          { id: 'val2', name: 'val2', meta: { type: 'number' } },
+          { id: 'val3', name: 'val3', meta: { type: 'number' } },
+          { id: 'split', name: 'split', meta: { type: 'string' } },
+        ],
+        rows: [
+          { val: 1, val2: 1, val3: 1, split: 'A' },
+          { val: 2, val2: 2, val3: 2, split: 'B' },
+          { val: 3, val2: 3, val3: 3, split: 'B' },
+          { val: 4, val2: 4, val3: 4, split: 'A' },
+          { val: 5, val2: 5, val3: 5, split: 'A' },
+          { val: 6, val2: 6, val3: 6, split: 'A' },
+          { val: 7, val2: 7, val3: 7, split: 'B' },
+          { val: 8, val2: 22, val3: 77, split: 'B' },
+        ],
+      },
+      { metric: ['val', 'val2', 'val3'], fn: ['sum'] }
+    );
+
+    expect(result.rows).toEqual([
+      {
+        val: 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8,
+        val2: 1 + 2 + 3 + 4 + 5 + 6 + 7 + 22,
+        val3: 1 + 2 + 3 + 4 + 5 + 6 + 7 + 77,
+      },
+    ]);
+  });
+
   it('can use different functions for each different metric', async () => {
     const result = await runFn(
       {
@@ -114,7 +147,7 @@ describe('collapse_fn', () => {
   };
 
   it('splits by a column', async () => {
-    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: 'sum' });
+    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: ['sum'] });
     expect(result.rows).toEqual([
       { val: 1 + 4 + 6, split: 'A' },
       { val: 2 + 7 + 8, split: 'B' },
@@ -123,7 +156,7 @@ describe('collapse_fn', () => {
   });
 
   it('applies avg', async () => {
-    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: 'avg' });
+    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: ['avg'] });
     expect(result.rows).toEqual([
       { val: (1 + 4 + 6) / 3, split: 'A' },
       { val: (2 + 7 + 8) / 3, split: 'B' },
@@ -132,7 +165,7 @@ describe('collapse_fn', () => {
   });
 
   it('applies min', async () => {
-    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: 'min' });
+    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: ['min'] });
     expect(result.rows).toEqual([
       { val: 1, split: 'A' },
       { val: 2, split: 'B' },
@@ -141,7 +174,7 @@ describe('collapse_fn', () => {
   });
 
   it('applies max', async () => {
-    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: 'max' });
+    const result = await runFn(twoSplitTable, { metric: ['val'], by: ['split'], fn: ['max'] });
     expect(result.rows).toEqual([
       { val: 6, split: 'A' },
       { val: 8, split: 'B' },
diff --git a/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.ts b/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.ts
index 5ca2248ed1ef7..ee3192705332d 100644
--- a/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.ts
+++ b/x-pack/plugins/lens/common/expressions/collapse/collapse_fn.ts
@@ -17,11 +17,8 @@ function getValueAsNumberArray(value: unknown) {
 }
 
 export const collapseFn: CollapseExpressionFunction['fn'] = (input, { by, metric, fn }) => {
-  const collapseFunctionsByMetricIndex = Array.isArray(fn)
-    ? fn
-    : metric
-    ? new Array(metric.length).fill(fn)
-    : [];
+  const collapseFunctionsByMetricIndex =
+    fn.length > 1 ? fn : metric ? new Array(metric.length).fill(fn[0]) : [];
 
   if (metric && metric.length !== collapseFunctionsByMetricIndex.length) {
     throw Error(`lens_collapse - Called with ${metric.length} metrics and ${fn.length} collapse functions. 
diff --git a/x-pack/plugins/lens/common/expressions/collapse/index.ts b/x-pack/plugins/lens/common/expressions/collapse/index.ts
index 5ea792e39cb0d..bd8df507c95e8 100644
--- a/x-pack/plugins/lens/common/expressions/collapse/index.ts
+++ b/x-pack/plugins/lens/common/expressions/collapse/index.ts
@@ -13,7 +13,7 @@ type CollapseFunction = 'sum' | 'avg' | 'min' | 'max';
 export interface CollapseArgs {
   by?: string[];
   metric?: string[];
-  fn: CollapseFunction | CollapseFunction[];
+  fn: CollapseFunction[];
 }
 
 /**