[Expressions] Refactor expression functions to use observables underneath (#100409) (#101988)

2025-04-23 09:19:04 -04:00 · 2021-06-11 12:00:14 +02:00 · 2021-06-11 12:00:14 +02:00 · f09704df30
commit f09704df30
parent 5c18bedd49
19 changed files with 604 additions and 453 deletions
--- a/docs/development/plugins/expressions/public/kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.md
+++ b/docs/development/plugins/expressions/public/kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.md
@ -23,7 +23,7 @@ export interface ExpressionFunctionDefinition<Name extends string, Input, Argume
 |  [help](./kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.help.md) | <code>string</code> | Help text displayed in the Expression editor. This text should be internationalized. |
 |  [inputTypes](./kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.inputtypes.md) | <code>Array&lt;TypeToString&lt;Input&gt;&gt;</code> | List of allowed type names for input value of this function. If this property is set the input of function will be cast to the first possible type in this list. If this property is missing the input will be provided to the function as-is. |
 |  [name](./kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.name.md) | <code>Name</code> | The name of the function, as will be used in expression. |
-|  [type](./kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.type.md) | <code>TypeToString&lt;UnwrapPromiseOrReturn&lt;Output&gt;&gt;</code> | Name of type of value this function outputs. |
+|  [type](./kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.type.md) | <code>TypeString&lt;Output&gt; &#124; UnmappedTypeStrings</code> | Name of type of value this function outputs. |

 ## Methods

--- a/docs/development/plugins/expressions/public/kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.type.md
+++ b/docs/development/plugins/expressions/public/kibana-plugin-plugins-expressions-public.expressionfunctiondefinition.type.md
@ -9,5 +9,5 @@ Name of type of value this function outputs.
 <b>Signature:</b>

 ```typescript
-type?: TypeToString<UnwrapPromiseOrReturn<Output>>;
+type?: TypeString<Output> | UnmappedTypeStrings;
 ```
--- a/docs/development/plugins/expressions/server/kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.md
+++ b/docs/development/plugins/expressions/server/kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.md
@ -23,7 +23,7 @@ export interface ExpressionFunctionDefinition<Name extends string, Input, Argume
 |  [help](./kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.help.md) | <code>string</code> | Help text displayed in the Expression editor. This text should be internationalized. |
 |  [inputTypes](./kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.inputtypes.md) | <code>Array&lt;TypeToString&lt;Input&gt;&gt;</code> | List of allowed type names for input value of this function. If this property is set the input of function will be cast to the first possible type in this list. If this property is missing the input will be provided to the function as-is. |
 |  [name](./kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.name.md) | <code>Name</code> | The name of the function, as will be used in expression. |
-|  [type](./kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.type.md) | <code>TypeToString&lt;UnwrapPromiseOrReturn&lt;Output&gt;&gt;</code> | Name of type of value this function outputs. |
+|  [type](./kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.type.md) | <code>TypeString&lt;Output&gt; &#124; UnmappedTypeStrings</code> | Name of type of value this function outputs. |

 ## Methods

--- a/docs/development/plugins/expressions/server/kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.type.md
+++ b/docs/development/plugins/expressions/server/kibana-plugin-plugins-expressions-server.expressionfunctiondefinition.type.md
@ -9,5 +9,5 @@ Name of type of value this function outputs.
 <b>Signature:</b>

 ```typescript
-type?: TypeToString<UnwrapPromiseOrReturn<Output>>;
+type?: TypeString<Output> | UnmappedTypeStrings;
 ```
--- a/src/plugins/expressions/common/expression_functions/specs/map_column.ts
+++ b/src/plugins/expressions/common/expression_functions/specs/map_column.ts
@ -6,8 +6,8 @@
 * Side Public License, v 1.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
+import { Observable, defer, of, zip } from 'rxjs';
+import { map } from 'rxjs/operators';
 import { i18n } from '@kbn/i18n';
 import { ExpressionFunctionDefinition } from '../types';
 import { Datatable, DatatableColumn, getType } from '../../expression_types';
@ -15,7 +15,7 @@ import { Datatable, DatatableColumn, getType } from '../../expression_types';
 export interface MapColumnArguments {
  id?: string | null;
  name: string;
-  expression?(datatable: Datatable): Observable<boolean | number | string | null>;
+  expression(datatable: Datatable): Observable<boolean | number | string | null>;
  copyMetaFrom?: string | null;
 }

@ -23,7 +23,7 @@ export const mapColumn: ExpressionFunctionDefinition<
  'mapColumn',
  Datatable,
  MapColumnArguments,
-  Promise<Datatable>
+  Observable<Datatable>
 > = {
  name: 'mapColumn',
  aliases: ['mc'], // midnight commander. So many times I've launched midnight commander instead of moving a file.
@ -80,57 +80,56 @@ export const mapColumn: ExpressionFunctionDefinition<
      default: null,
    },
  },
-  fn: (input, args) => {
-    const expression = (...params: Parameters<Required<MapColumnArguments>['expression']>) =>
-      args
-        .expression?.(...params)
-        .pipe(take(1))
-        .toPromise() ?? Promise.resolve(null);
+  fn(input, args) {
+    const existingColumnIndex = input.columns.findIndex(({ id, name }) =>
+      args.id ? id === args.id : name === args.name
+    );
+    const id = input.columns[existingColumnIndex]?.id ?? args.id ?? args.name;

-    const columns = [...input.columns];
-    const existingColumnIndex = columns.findIndex(({ id, name }) => {
-      if (args.id) {
-        return id === args.id;
-      }
-      return name === args.name;
-    });
-    const columnId =
-      existingColumnIndex === -1 ? args.id ?? args.name : columns[existingColumnIndex].id;
+    return defer(() => {
+      const rows$ = input.rows.length
+        ? zip(
+            ...input.rows.map((row) =>
+              args
+                .expression({
+                  type: 'datatable',
+                  columns: [...input.columns],
+                  rows: [row],
+                })
+                .pipe(map((value) => ({ ...row, [id]: value })))
+            )
+          )
+        : of([]);

-    const rowPromises = input.rows.map((row) => {
-      return expression({
-        type: 'datatable',
-        columns,
-        rows: [row],
-      }).then((val) => ({
-        ...row,
-        [columnId]: val,
-      }));
-    });
+      return rows$.pipe<Datatable>(
+        map((rows) => {
+          const type = getType(rows[0]?.[id]);
+          const newColumn: DatatableColumn = {
+            id,
+            name: args.name,
+            meta: { type, params: { id: type } },
+          };
+          if (args.copyMetaFrom) {
+            const metaSourceFrom = input.columns.find(
+              ({ id: columnId }) => columnId === args.copyMetaFrom
+            );
+            newColumn.meta = { ...newColumn.meta, ...(metaSourceFrom?.meta ?? {}) };
+          }

-    return Promise.all(rowPromises).then((rows) => {
-      const type = rows.length ? getType(rows[0][columnId]) : 'null';
-      const newColumn: DatatableColumn = {
-        id: columnId,
-        name: args.name,
-        meta: { type, params: { id: type } },
-      };
-      if (args.copyMetaFrom) {
-        const metaSourceFrom = columns.find(({ id }) => id === args.copyMetaFrom);
-        newColumn.meta = { ...newColumn.meta, ...(metaSourceFrom?.meta || {}) };
-      }
+          const columns = [...input.columns];
+          if (existingColumnIndex === -1) {
+            columns.push(newColumn);
+          } else {
+            columns[existingColumnIndex] = newColumn;
+          }

-      if (existingColumnIndex === -1) {
-        columns.push(newColumn);
-      } else {
-        columns[existingColumnIndex] = newColumn;
-      }
-
-      return {
-        type: 'datatable',
-        columns,
-        rows,
-      } as Datatable;
+          return {
+            columns,
+            rows,
+            type: 'datatable',
+          };
+        })
+      );
    });
  },
 };
--- a/src/plugins/expressions/common/expression_functions/specs/tests/map_column.test.ts
+++ b/src/plugins/expressions/common/expression_functions/specs/tests/map_column.test.ts
@ -6,7 +6,8 @@
 * Side Public License, v 1.
 */

-import { of } from 'rxjs';
+import { of, Observable } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { Datatable } from '../../../expression_types';
 import { mapColumn, MapColumnArguments } from '../map_column';
 import { emptyTable, functionWrapper, testTable } from './utils';
@ -16,142 +17,227 @@ const pricePlusTwo = (datatable: Datatable) => of(datatable.rows[0].price + 2);
 describe('mapColumn', () => {
  const fn = functionWrapper(mapColumn);
  const runFn = (input: Datatable, args: MapColumnArguments) =>
-    fn(input, args) as Promise<Datatable>;
+    fn(input, args) as Observable<Datatable>;
+  let testScheduler: TestScheduler;

-  it('returns a datatable with a new column with the values from mapping a function over each row in a datatable', async () => {
-    const arbitraryRowIndex = 2;
-    const result = await runFn(testTable, {
-      id: 'pricePlusTwo',
-      name: 'pricePlusTwo',
-      expression: pricePlusTwo,
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });
+
+  it('returns a datatable with a new column with the values from mapping a function over each row in a datatable', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(testTable, {
+          id: 'pricePlusTwo',
+          name: 'pricePlusTwo',
+          expression: pricePlusTwo,
+        })
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: [
+            ...testTable.columns,
+            {
+              id: 'pricePlusTwo',
+              name: 'pricePlusTwo',
+              meta: { type: 'number', params: { id: 'number' } },
+            },
+          ],
+          rows: expect.arrayContaining([
+            expect.objectContaining({
+              pricePlusTwo: expect.anything(),
+            }),
+          ]),
+        }),
+      ]);
    });
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toEqual([
-      ...testTable.columns,
-      {
-        id: 'pricePlusTwo',
-        name: 'pricePlusTwo',
-        meta: { type: 'number', params: { id: 'number' } },
-      },
-    ]);
-    expect(result.columns[result.columns.length - 1]).toHaveProperty('name', 'pricePlusTwo');
-    expect(result.rows[arbitraryRowIndex]).toHaveProperty('pricePlusTwo');
  });

-  it('allows the id arg to be optional, looking up by name instead', async () => {
-    const result = await runFn(testTable, { name: 'name label', expression: pricePlusTwo });
-    const nameColumnIndex = result.columns.findIndex(({ name }) => name === 'name label');
-    const arbitraryRowIndex = 4;
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(testTable.columns.length);
-    expect(result.columns[nameColumnIndex]).toHaveProperty('id', 'name');
-    expect(result.columns[nameColumnIndex]).toHaveProperty('name', 'name label');
-    expect(result.columns[nameColumnIndex].meta).toHaveProperty('type', 'number');
-    expect(result.rows[arbitraryRowIndex]).toHaveProperty('name', 202);
-    expect(result.rows[arbitraryRowIndex]).not.toHaveProperty('name label');
-  });
-
-  it('allows a duplicate name when the ids are different', async () => {
-    const result = await runFn(testTable, {
-      id: 'new',
-      name: 'name label',
-      expression: pricePlusTwo,
+  it('allows the id arg to be optional, looking up by name instead', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(runFn(testTable, { name: 'name label', expression: pricePlusTwo })).toBe(
+        '(0|)',
+        [
+          expect.objectContaining({
+            type: 'datatable',
+            columns: expect.arrayContaining([
+              expect.objectContaining({
+                id: 'name',
+                name: 'name label',
+                meta: expect.objectContaining({ type: 'number' }),
+              }),
+            ]),
+            rows: expect.arrayContaining([
+              expect.objectContaining({
+                name: 202,
+              }),
+            ]),
+          }),
+        ]
+      );
    });
-    const nameColumnIndex = result.columns.findIndex(({ id }) => id === 'new');
-    const arbitraryRowIndex = 4;
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(testTable.columns.length + 1);
-    expect(result.columns[nameColumnIndex]).toHaveProperty('id', 'new');
-    expect(result.columns[nameColumnIndex]).toHaveProperty('name', 'name label');
-    expect(result.columns[nameColumnIndex].meta).toHaveProperty('type', 'number');
-    expect(result.rows[arbitraryRowIndex]).toHaveProperty('new', 202);
  });

-  it('adds a column to empty tables', async () => {
-    const result = await runFn(emptyTable, { name: 'name', expression: pricePlusTwo });
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(1);
-    expect(result.columns[0]).toHaveProperty('name', 'name');
-    expect(result.columns[0].meta).toHaveProperty('type', 'null');
+  it('allows a duplicate name when the ids are different', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(testTable, {
+          id: 'new',
+          name: 'name label',
+          expression: pricePlusTwo,
+        })
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: expect.arrayContaining([
+            expect.objectContaining({
+              id: 'new',
+              name: 'name label',
+              meta: expect.objectContaining({ type: 'number' }),
+            }),
+          ]),
+          rows: expect.arrayContaining([
+            expect.objectContaining({
+              new: 202,
+            }),
+          ]),
+        }),
+      ]);
+    });
  });

-  it('should assign specific id, different from name, when id arg is passed for new columns', async () => {
-    const result = await runFn(emptyTable, { name: 'name', id: 'myid', expression: pricePlusTwo });
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(1);
-    expect(result.columns[0]).toHaveProperty('name', 'name');
-    expect(result.columns[0]).toHaveProperty('id', 'myid');
-    expect(result.columns[0].meta).toHaveProperty('type', 'null');
+  it('overwrites existing column with the new column if an existing column name is provided', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(runFn(testTable, { name: 'name', expression: pricePlusTwo })).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: expect.arrayContaining([
+            expect.objectContaining({
+              name: 'name',
+              meta: expect.objectContaining({ type: 'number' }),
+            }),
+          ]),
+          rows: expect.arrayContaining([
+            expect.objectContaining({
+              name: 202,
+            }),
+          ]),
+        }),
+      ]);
+    });
  });

-  it('should copy over the meta information from the specified column', async () => {
-    const result = await runFn(
-      {
-        ...testTable,
-        columns: [
-          ...testTable.columns,
-          // add a new entry
+  it('adds a column to empty tables', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(runFn(emptyTable, { name: 'name', expression: pricePlusTwo })).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: [
+            expect.objectContaining({
+              name: 'name',
+              meta: expect.objectContaining({ type: 'null' }),
+            }),
+          ],
+        }),
+      ]);
+    });
+  });
+
+  it('should assign specific id, different from name, when id arg is passed for copied column', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(testTable, { name: 'name', id: 'myid', expression: pricePlusTwo })
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: expect.arrayContaining([
+            expect.objectContaining({
+              id: 'myid',
+              name: 'name',
+              meta: expect.objectContaining({ type: 'number' }),
+            }),
+          ]),
+        }),
+      ]);
+    });
+  });
+
+  it('should copy over the meta information from the specified column', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(
          {
-            id: 'myId',
-            name: 'myName',
-            meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+            ...testTable,
+            columns: [
+              ...testTable.columns,
+              // add a new entry
+              {
+                id: 'myId',
+                name: 'myName',
+                meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+              },
+            ],
+            rows: testTable.rows.map((row) => ({ ...row, myId: Date.now() })),
          },
-        ],
-        rows: testTable.rows.map((row) => ({ ...row, myId: Date.now() })),
-      },
-      { name: 'name', copyMetaFrom: 'myId', expression: pricePlusTwo }
-    );
-    const nameColumnIndex = result.columns.findIndex(({ name }) => name === 'name');
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns[nameColumnIndex]).toEqual({
-      id: 'name',
-      name: 'name',
-      meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+          { name: 'name', copyMetaFrom: 'myId', expression: pricePlusTwo }
+        )
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: expect.arrayContaining([
+            expect.objectContaining({
+              id: 'name',
+              name: 'name',
+              meta: { type: 'date', params: { id: 'number', params: { digits: 2 } } },
+            }),
+          ]),
+        }),
+      ]);
    });
  });

-  it('should be resilient if the references column for meta information does not exists', async () => {
-    const result = await runFn(emptyTable, {
-      name: 'name',
-      copyMetaFrom: 'time',
-      expression: pricePlusTwo,
+  it('should be resilient if the references column for meta information does not exists', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(emptyTable, {
+          name: 'name',
+          copyMetaFrom: 'time',
+          expression: pricePlusTwo,
+        })
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: [
+            expect.objectContaining({
+              id: 'name',
+              name: 'name',
+              meta: expect.objectContaining({ type: 'null' }),
+            }),
+          ],
+        }),
+      ]);
    });
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(1);
-    expect(result.columns[0]).toHaveProperty('name', 'name');
-    expect(result.columns[0]).toHaveProperty('id', 'name');
-    expect(result.columns[0].meta).toHaveProperty('type', 'null');
  });

-  it('should correctly infer the type fromt he first row if the references column for meta information does not exists', async () => {
-    const result = await runFn(
-      { ...emptyTable, rows: [...emptyTable.rows, { value: 5 }] },
-      { name: 'value', copyMetaFrom: 'time', expression: pricePlusTwo }
-    );
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toHaveLength(1);
-    expect(result.columns[0]).toHaveProperty('name', 'value');
-    expect(result.columns[0]).toHaveProperty('id', 'value');
-    expect(result.columns[0].meta).toHaveProperty('type', 'number');
-  });
-
-  describe('expression', () => {
-    it('maps null values to the new column', async () => {
-      const result = await runFn(testTable, { name: 'empty' });
-      const emptyColumnIndex = result.columns.findIndex(({ name }) => name === 'empty');
-      const arbitraryRowIndex = 8;
-
-      expect(result.columns[emptyColumnIndex]).toHaveProperty('name', 'empty');
-      expect(result.columns[emptyColumnIndex].meta).toHaveProperty('type', 'null');
-      expect(result.rows[arbitraryRowIndex]).toHaveProperty('empty', null);
+  it('should correctly infer the type fromt he first row if the references column for meta information does not exists', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        runFn(
+          { ...emptyTable, rows: [...emptyTable.rows, { value: 5 }] },
+          { name: 'value', copyMetaFrom: 'time', expression: pricePlusTwo }
+        )
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: [
+            expect.objectContaining({
+              id: 'value',
+              name: 'value',
+              meta: expect.objectContaining({ type: 'number' }),
+            }),
+          ],
+        }),
+      ]);
    });
  });
 });
--- a/src/plugins/expressions/common/expression_functions/types.ts
+++ b/src/plugins/expressions/common/expression_functions/types.ts
@ -6,9 +6,8 @@
 * Side Public License, v 1.
 */

-import { UnwrapPromiseOrReturn } from '@kbn/utility-types';
 import { ArgumentType } from './arguments';
-import { TypeToString } from '../types/common';
+import { TypeToString, TypeString, UnmappedTypeStrings } from '../types/common';
 import { ExecutionContext } from '../execution/types';
 import {
  ExpressionFunctionClog,
@ -47,7 +46,7 @@ export interface ExpressionFunctionDefinition<
  /**
   * Name of type of value this function outputs.
   */
-  type?: TypeToString<UnwrapPromiseOrReturn<Output>>;
+  type?: TypeString<Output> | UnmappedTypeStrings;

  /**
   * List of allowed type names for input value of this function. If this
--- a/src/plugins/expressions/public/public.api.md
+++ b/src/plugins/expressions/public/public.api.md
@ -375,7 +375,7 @@ export interface ExpressionFunctionDefinition<Name extends string, Input, Argume
    help: string;
    inputTypes?: Array<TypeToString<Input>>;
    name: Name;
-    type?: TypeToString<UnwrapPromiseOrReturn<Output>>;
+    type?: TypeString<Output> | UnmappedTypeStrings;
 }

 // @public
--- a/src/plugins/expressions/server/server.api.md
+++ b/src/plugins/expressions/server/server.api.md
@ -347,7 +347,7 @@ export interface ExpressionFunctionDefinition<Name extends string, Input, Argume
    help: string;
    inputTypes?: Array<TypeToString<Input>>;
    name: Name;
-    type?: TypeToString<UnwrapPromiseOrReturn<Output>>;
+    type?: TypeString<Output> | UnmappedTypeStrings;
 }

 // @public
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/case.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/case.test.js
@ -6,11 +6,17 @@
 */

 import { of } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { functionWrapper } from '../../../test_helpers/function_wrapper';
 import { caseFn } from './case';

 describe('case', () => {
  const fn = functionWrapper(caseFn);
+  let testScheduler;
+
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });

  describe('spec', () => {
    it('is a function', () => {
@ -19,29 +25,22 @@ describe('case', () => {
  });

  describe('function', () => {
-    describe('no args', () => {
-      it('should return a case object that matches with the result as the context', () => {
-        const context = null;
-        const args = {};
-        expect(fn(context, args)).resolves.toEqual({
-          type: 'case',
-          matches: true,
-          result: context,
-        });
-      });
-    });
-
    describe('no if or value', () => {
      it('should return the result if provided', () => {
        const context = null;
        const args = {
          then: () => of('foo'),
        };
-        expect(fn(context, args)).resolves.toEqual({
-          type: 'case',
-          matches: true,
-          result: 'foo',
-        });
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', [
+            {
+              type: 'case',
+              matches: true,
+              result: 'foo',
+            },
+          ])
+        );
      });
    });

@ -49,11 +48,16 @@ describe('case', () => {
      it('should return as the matches prop', () => {
        const context = null;
        const args = { if: false };
-        expect(fn(context, args)).resolves.toEqual({
-          type: 'case',
-          matches: args.if,
-          result: context,
-        });
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', [
+            {
+              type: 'case',
+              matches: args.if,
+              result: context,
+            },
+          ])
+        );
      });
    });

@ -63,15 +67,23 @@ describe('case', () => {
          when: () => of('foo'),
          then: () => of('bar'),
        };
-        expect(fn('foo', args)).resolves.toEqual({
-          type: 'case',
-          matches: true,
-          result: 'bar',
-        });
-        expect(fn('bar', args)).resolves.toEqual({
-          type: 'case',
-          matches: false,
-          result: null,
+
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(fn('foo', args)).toBe('(0|)', [
+            {
+              type: 'case',
+              matches: true,
+              result: 'bar',
+            },
+          ]);
+
+          expectObservable(fn('bar', args)).toBe('(0|)', [
+            {
+              type: 'case',
+              matches: false,
+              result: null,
+            },
+          ]);
        });
      });
    });
@ -81,13 +93,18 @@ describe('case', () => {
        const context = null;
        const args = {
          when: () => 'foo',
-          if: true,
+          if: false,
        };
-        expect(fn(context, args)).resolves.toEqual({
-          type: 'case',
-          matches: args.if,
-          result: context,
-        });
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', [
+            {
+              type: 'case',
+              matches: args.if,
+              result: context,
+            },
+          ])
+        );
      });
    });
  });
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/case.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/case.ts
@ -5,15 +5,15 @@
 * 2.0.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
+import { Observable, defer, isObservable, of } from 'rxjs';
+import { map, concatMap } from 'rxjs/operators';
 import { ExpressionFunctionDefinition } from 'src/plugins/expressions/common';
 import { getFunctionHelp } from '../../../i18n';

 interface Arguments {
  when?(): Observable<any>;
  if?: boolean;
-  then?(): Observable<any>;
+  then(): Observable<any>;
 }

 interface Case {
@ -22,7 +22,7 @@ interface Case {
  result: any;
 }

-export function caseFn(): ExpressionFunctionDefinition<'case', any, Arguments, Promise<Case>> {
+export function caseFn(): ExpressionFunctionDefinition<'case', any, Arguments, Observable<Case>> {
  const { help, args: argHelp } = getFunctionHelp().case;

  return {
@ -45,24 +45,24 @@ export function caseFn(): ExpressionFunctionDefinition<'case', any, Arguments, P
        help: argHelp.then!,
      },
    },
-    fn: async (input, args) => {
-      const matches = await doesMatch(input, args);
-      const result = matches ? await getResult(input, args) : null;
-      return { type: 'case', matches, result };
+    fn(input, { if: condition, then, when }) {
+      return defer(() => {
+        const matches = condition ?? when?.().pipe(map((value) => value === input)) ?? true;
+
+        return isObservable(matches) ? matches : of(matches);
+      }).pipe(
+        concatMap((matches) =>
+          (matches ? then() : of(null)).pipe(
+            map(
+              (result): Case => ({
+                matches,
+                result,
+                type: 'case',
+              })
+            )
+          )
+        )
+      );
    },
  };
 }
-
-async function doesMatch(context: any, args: Arguments) {
-  if (typeof args.if !== 'undefined') {
-    return args.if;
-  }
-  if (typeof args.when !== 'undefined') {
-    return (await args.when().pipe(take(1)).toPromise()) === context;
-  }
-  return true;
-}
-
-async function getResult(context: any, args: Arguments) {
-  return args.then?.().pipe(take(1)).toPromise() ?? context;
-}
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/filterrows.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/filterrows.test.js
@ -6,6 +6,7 @@
 */

 import { of } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { functionWrapper } from '../../../test_helpers/function_wrapper';
 import { testTable } from './__fixtures__/test_tables';
 import { filterrows } from './filterrows';
@ -15,31 +16,46 @@ const returnFalse = () => of(false);

 describe('filterrows', () => {
  const fn = functionWrapper(filterrows);
+  let testScheduler;
+
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });

  it('returns a datable', () => {
-    expect(fn(testTable, { fn: inStock })).resolves.toHaveProperty('type', 'datatable');
+    testScheduler.run(({ expectObservable }) =>
+      expectObservable(fn(testTable, { fn: inStock })).toBe('(0|)', [
+        expect.objectContaining({ type: 'datatable' }),
+      ])
+    );
  });

  it('keeps rows that evaluate to true and removes rows that evaluate to false', () => {
    const inStockRows = testTable.rows.filter((row) => row.in_stock);

-    expect(fn(testTable, { fn: inStock })).resolves.toEqual(
-      expect.objectContaining({
-        columns: testTable.columns,
-        rows: inStockRows,
-      })
+    testScheduler.run(({ expectObservable }) =>
+      expectObservable(fn(testTable, { fn: inStock })).toBe('(0|)', [
+        expect.objectContaining({
+          columns: testTable.columns,
+          rows: inStockRows,
+        }),
+      ])
    );
  });

  it('returns datatable with no rows when no rows meet function condition', () => {
-    expect(fn(testTable, { fn: returnFalse })).resolves.toEqual(
-      expect.objectContaining({
-        rows: [],
-      })
+    testScheduler.run(({ expectObservable }) =>
+      expectObservable(fn(testTable, { fn: returnFalse })).toBe('(0|)', [
+        expect.objectContaining({
+          rows: [],
+        }),
+      ])
    );
  });

  it('throws when no function is provided', () => {
-    expect(() => fn(testTable)).toThrow('fn is not a function');
+    testScheduler.run(({ expectObservable }) =>
+      expectObservable(fn(testTable)).toBe('#', {}, new TypeError('fn is not a function'))
+    );
  });
 });
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/filterrows.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/filterrows.ts
@ -5,8 +5,8 @@
 * 2.0.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
+import { Observable, combineLatest, defer } from 'rxjs';
+import { map } from 'rxjs/operators';
 import { Datatable, ExpressionFunctionDefinition } from '../../../types';
 import { getFunctionHelp } from '../../../i18n';

@ -18,7 +18,7 @@ export function filterrows(): ExpressionFunctionDefinition<
  'filterrows',
  Datatable,
  Arguments,
-  Promise<Datatable>
+  Observable<Datatable>
 > {
  const { help, args: argHelp } = getFunctionHelp().filterrows;

@ -38,24 +38,12 @@ export function filterrows(): ExpressionFunctionDefinition<
      },
    },
    fn(input, { fn }) {
-      const checks = input.rows.map((row) =>
-        fn({
-          ...input,
-          rows: [row],
-        })
-          .pipe(take(1))
-          .toPromise()
+      return defer(() =>
+        combineLatest(input.rows.map((row) => fn({ ...input, rows: [row] })))
+      ).pipe(
+        map((checks) => input.rows.filter((row, i) => checks[i])),
+        map((rows) => ({ ...input, rows }))
      );
-
-      return Promise.all(checks)
-        .then((results) => input.rows.filter((row, i) => results[i]))
-        .then(
-          (rows) =>
-            ({
-              ...input,
-              rows,
-            } as Datatable)
-        );
    },
  };
 }
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/if.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/if.test.js
@ -6,11 +6,17 @@
 */

 import { of } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { functionWrapper } from '../../../test_helpers/function_wrapper';
 import { ifFn } from './if';

 describe('if', () => {
  const fn = functionWrapper(ifFn);
+  let testScheduler;
+
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });

  describe('spec', () => {
    it('is a function', () => {
@ -21,66 +27,73 @@ describe('if', () => {
  describe('function', () => {
    describe('condition passed', () => {
      it('with then', () => {
-        expect(fn(null, { condition: true, then: () => of('foo') })).resolves.toBe('foo');
-        expect(
-          fn(null, { condition: true, then: () => of('foo'), else: () => of('bar') })
-        ).resolves.toBe('foo');
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(fn(null, { condition: true, then: () => of('foo') })).toBe('(0|)', [
+            'foo',
+          ]);
+
+          expectObservable(
+            fn(null, { condition: true, then: () => of('foo'), else: () => of('bar') })
+          ).toBe('(0|)', ['foo']);
+        });
      });

      it('without then', () => {
-        expect(fn(null, { condition: true })).resolves.toBe(null);
-        expect(fn('some context', { condition: true })).resolves.toBe('some context');
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(fn(null, { condition: true })).toBe('(0|)', [null]);
+          expectObservable(fn('some context', { condition: true })).toBe('(0|)', ['some context']);
+        });
      });
    });

    describe('condition failed', () => {
-      it('with else', () =>
-        expect(
-          fn('some context', {
-            condition: false,
-            then: () => of('foo'),
-            else: () => of('bar'),
-          })
-        ).resolves.toBe('bar'));
-
-      it('without else', () =>
-        expect(fn('some context', { condition: false, then: () => of('foo') })).resolves.toBe(
-          'some context'
-        ));
-    });
-
-    describe('falsy values', () => {
-      describe('for then', () => {
-        it('with null', () => {
-          expect(fn('some context', { condition: true, then: () => of(null) })).resolves.toBe(null);
-        });
-
-        it('with false', () => {
-          expect(fn('some context', { condition: true, then: () => of(false) })).resolves.toBe(
-            false
-          );
-        });
-
-        it('with 0', () => {
-          expect(fn('some context', { condition: true, then: () => of(0) })).resolves.toBe(0);
+      it('with else', () => {
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(
+            fn('some context', {
+              condition: false,
+              then: () => of('foo'),
+              else: () => of('bar'),
+            })
+          ).toBe('(0|)', ['bar']);
        });
      });

-      describe('for else', () => {
-        it('with null', () => {
-          expect(fn('some context', { condition: false, else: () => of(null) })).resolves.toBe(
-            null
-          );
+      it('without else', () => {
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(
+            fn('some context', { condition: false, then: () => of('foo') })
+          ).toBe('(0|)', ['some context']);
        });
+      });
+    });

-        it('with false', () => {
-          expect(fn('some context', { condition: false, else: () => of(false) })).resolves.toBe(
-            false
-          );
+    describe('falsy values', () => {
+      // eslint-disable-next-line no-unsanitized/method
+      it.each`
+        value
+        ${null}
+        ${false}
+        ${0}
+      `('for then with $value', ({ value }) => {
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(
+            fn('some context', { condition: true, then: () => of(value) })
+          ).toBe('(0|)', [value]);
        });
+      });

-        it('with 0', () => {
-          expect(fn('some context', { condition: false, else: () => of(0) })).resolves.toBe(0);
+      // eslint-disable-next-line no-unsanitized/method
+      it.each`
+        value
+        ${null}
+        ${false}
+        ${0}
+      `('for else with $value', ({ value }) => {
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(
+            fn('some context', { condition: false, else: () => of(value) })
+          ).toBe('(0|)', [value]);
        });
      });
    });
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/if.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/if.ts
@ -5,18 +5,22 @@
 * 2.0.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
+import { Observable, defer, of } from 'rxjs';
 import { ExpressionFunctionDefinition } from 'src/plugins/expressions/common';
 import { getFunctionHelp } from '../../../i18n';

 interface Arguments {
-  condition: boolean | null;
+  condition: boolean;
  then?(): Observable<any>;
  else?(): Observable<any>;
 }

-export function ifFn(): ExpressionFunctionDefinition<'if', unknown, Arguments, unknown> {
+export function ifFn(): ExpressionFunctionDefinition<
+  'if',
+  unknown,
+  Arguments,
+  Observable<unknown>
+> {
  const { help, args: argHelp } = getFunctionHelp().if;

  return {
@ -38,12 +42,8 @@ export function ifFn(): ExpressionFunctionDefinition<'if', unknown, Arguments, u
        help: argHelp.else!,
      },
    },
-    fn: async (input, args) => {
-      if (args.condition) {
-        return args.then?.().pipe(take(1)).toPromise() ?? input;
-      } else {
-        return args.else?.().pipe(take(1)).toPromise() ?? input;
-      }
+    fn(input, args) {
+      return defer(() => (args.condition ? args.then?.() : args.else?.()) ?? of(input));
    },
  };
 }
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/ply.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/ply.test.js
@ -6,6 +6,7 @@
 */

 import { of } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { functionWrapper } from '../../../test_helpers/function_wrapper';
 import { getFunctionErrors } from '../../../i18n';
 import { testTable } from './__fixtures__/test_tables';
@ -46,36 +47,56 @@ const rowCount = (datatable) =>

 describe('ply', () => {
  const fn = functionWrapper(ply);
+  let testScheduler;

-  it('maps a function over sub datatables grouped by specified columns and merges results into one datatable', async () => {
-    const arbitaryRowIndex = 0;
-    const result = await fn(testTable, {
-      by: ['name', 'in_stock'],
-      expression: [averagePrice, rowCount],
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });
+
+  it('maps a function over sub datatables grouped by specified columns and merges results into one datatable', () => {
+    testScheduler.run(({ expectObservable }) => {
+      expectObservable(
+        fn(testTable, {
+          by: ['name', 'in_stock'],
+          expression: [averagePrice, rowCount],
+        })
+      ).toBe('(0|)', [
+        expect.objectContaining({
+          type: 'datatable',
+          columns: [
+            { id: 'name', name: 'name', meta: { type: 'string' } },
+            { id: 'in_stock', name: 'in_stock', meta: { type: 'boolean' } },
+            { id: 'average_price', name: 'average_price', meta: { type: 'number' } },
+            { id: 'row_count', name: 'row_count', meta: { type: 'number' } },
+          ],
+          rows: expect.arrayContaining([
+            expect.objectContaining({
+              average_price: expect.anything(),
+              row_count: expect.anything(),
+            }),
+          ]),
+        }),
+      ]);
    });
-
-    expect(result.type).toBe('datatable');
-    expect(result.columns).toEqual([
-      { id: 'name', name: 'name', meta: { type: 'string' } },
-      { id: 'in_stock', name: 'in_stock', meta: { type: 'boolean' } },
-      { id: 'average_price', name: 'average_price', meta: { type: 'number' } },
-      { id: 'row_count', name: 'row_count', meta: { type: 'number' } },
-    ]);
-    expect(result.rows[arbitaryRowIndex]).toHaveProperty('average_price');
-    expect(result.rows[arbitaryRowIndex]).toHaveProperty('row_count');
  });

  describe('missing args', () => {
    it('returns the original datatable if both args are missing', () => {
-      expect(fn(testTable)).resolves.toEqual(testTable);
+      testScheduler.run(({ expectObservable }) => {
+        expectObservable(fn(testTable)).toBe('(0|)', [testTable]);
+      });
    });

    describe('by', () => {
      it('passes the entire context into the expression when no columns are provided', () => {
-        expect(fn(testTable, { expression: [rowCount] })).resolves.toEqual({
-          type: 'datatable',
-          rows: [{ row_count: testTable.rows.length }],
-          columns: [{ id: 'row_count', name: 'row_count', meta: { type: 'number' } }],
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(fn(testTable, { expression: [rowCount] })).toBe('(0|)', [
+            {
+              type: 'datatable',
+              rows: [{ row_count: testTable.rows.length }],
+              columns: [{ id: 'row_count', name: 'row_count', meta: { type: 'number' } }],
+            },
+          ]);
        });
      });

@ -91,24 +112,37 @@ describe('ply', () => {
    });

    describe('expression', () => {
-      it('returns the original datatable grouped by the specified columns', async () => {
-        const arbitaryRowIndex = 6;
-        const result = await fn(testTable, { by: ['price', 'quantity'] });
-
-        expect(result.columns[0]).toHaveProperty('name', 'price');
-        expect(result.columns[1]).toHaveProperty('name', 'quantity');
-        expect(result.rows[arbitaryRowIndex]).toHaveProperty('price');
-        expect(result.rows[arbitaryRowIndex]).toHaveProperty('quantity');
+      it('returns the original datatable grouped by the specified columns', () => {
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(fn(testTable, { by: ['price', 'quantity'] })).toBe('(0|)', [
+            expect.objectContaining({
+              columns: expect.arrayContaining([
+                expect.objectContaining({ name: 'price' }),
+                expect.objectContaining({ name: 'quantity' }),
+              ]),
+              rows: expect.arrayContaining([
+                expect.objectContaining({
+                  price: expect.anything(),
+                  quantity: expect.anything(),
+                }),
+              ]),
+            }),
+          ]);
+        });
      });

      it('throws when row counts do not match across resulting datatables', () => {
-        expect(
-          fn(testTable, { by: ['name'], expression: [doublePrice, rowCount] })
-        ).rejects.toEqual(
-          expect.objectContaining({
-            message: errors.rowCountMismatch().message,
-          })
-        );
+        testScheduler.run(({ expectObservable }) => {
+          expectObservable(
+            fn(testTable, { by: ['name'], expression: [doublePrice, rowCount] })
+          ).toBe(
+            '#',
+            [],
+            expect.objectContaining({
+              message: errors.rowCountMismatch().message,
+            })
+          );
+        });
      });
    });
  });
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/ply.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/ply.ts
@ -5,18 +5,18 @@
 * 2.0.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
-import { groupBy, flatten, pick, map } from 'lodash';
+import { combineLatest, defer, of, Observable } from 'rxjs';
+import { map } from 'rxjs/operators';
+import { groupBy, flatten, pick, map as _map, uniqWith } from 'lodash';
 import { Datatable, DatatableColumn, ExpressionFunctionDefinition } from '../../../types';
 import { getFunctionHelp, getFunctionErrors } from '../../../i18n';

 interface Arguments {
-  by: string[];
+  by?: string[];
  expression: Array<(datatable: Datatable) => Observable<Datatable>>;
 }

-type Output = Datatable | Promise<Datatable>;
+type Output = Datatable | Observable<Datatable>;

 export function ply(): ExpressionFunctionDefinition<'ply', Datatable, Arguments, Output> {
  const { help, args: argHelp } = getFunctionHelp().ply;
@ -30,7 +30,7 @@ export function ply(): ExpressionFunctionDefinition<'ply', Datatable, Arguments,
    args: {
      by: {
        types: ['string'],
-        help: argHelp.by,
+        help: argHelp.by!,
        multi: true,
      },
      expression: {
@ -41,86 +41,62 @@ export function ply(): ExpressionFunctionDefinition<'ply', Datatable, Arguments,
        help: argHelp.expression,
      },
    },
-    fn: (input, args) => {
+    fn(input, args) {
      if (!args) {
        return input;
      }

-      let byColumns: DatatableColumn[];
-      let originalDatatables: Datatable[];
-
-      if (args.by) {
-        byColumns = args.by.map((by) => {
-          const column = input.columns.find((col) => col.name === by);
+      const byColumns =
+        args.by?.map((by) => {
+          const column = input.columns.find(({ name }) => name === by);

          if (!column) {
            throw errors.columnNotFound(by);
          }

          return column;
-        });
+        }) ?? [];

-        const keyedDatatables = groupBy(input.rows, (row) => JSON.stringify(pick(row, args.by)));
+      const originalDatatables = args.by
+        ? Object.values(
+            groupBy(input.rows, (row) => JSON.stringify(pick(row, args.by!)))
+          ).map((rows) => ({ ...input, rows }))
+        : [input];

-        originalDatatables = Object.values(keyedDatatables).map((rows) => ({
-          ...input,
-          rows,
-        }));
-      } else {
-        originalDatatables = [input];
-      }
+      const datatables$ = originalDatatables.map((originalDatatable) =>
+        combineLatest(
+          args.expression?.map((expression) => defer(() => expression(originalDatatable))) ?? [
+            of(originalDatatable),
+          ]
+        ).pipe(map(combineAcross))
+      );

-      const datatablePromises = originalDatatables.map((originalDatatable) => {
-        let expressionResultPromises = [];
-
-        if (args.expression) {
-          expressionResultPromises = args.expression.map((expression) =>
-            expression(originalDatatable).pipe(take(1)).toPromise()
+      return (datatables$.length ? combineLatest(datatables$) : of([])).pipe(
+        map((newDatatables) => {
+          // Here we're just merging each for the by splits, so it doesn't actually matter if the rows are the same length
+          const columns = combineColumns([byColumns].concat(_map(newDatatables, 'columns')));
+          const rows = flatten(
+            newDatatables.map((datatable, index) =>
+              datatable.rows.map((row) => ({
+                ...pick(originalDatatables[index].rows[0], args.by!),
+                ...row,
+              }))
+            )
          );
-        } else {
-          expressionResultPromises.push(Promise.resolve(originalDatatable));
-        }

-        return Promise.all(expressionResultPromises).then(combineAcross);
-      });
-
-      return Promise.all(datatablePromises).then((newDatatables) => {
-        // Here we're just merging each for the by splits, so it doesn't actually matter if the rows are the same length
-        const columns = combineColumns([byColumns].concat(map(newDatatables, 'columns')));
-        const rows = flatten(
-          newDatatables.map((dt, i) => {
-            const byColumnValues = pick(originalDatatables[i].rows[0], args.by);
-            return dt.rows.map((row) => ({
-              ...byColumnValues,
-              ...row,
-            }));
-          })
-        );
-
-        return {
-          type: 'datatable',
-          rows,
-          columns,
-        } as Datatable;
-      });
+          return {
+            type: 'datatable',
+            rows,
+            columns,
+          } as Datatable;
+        })
+      );
    },
  };
 }

 function combineColumns(arrayOfColumnsArrays: DatatableColumn[][]) {
-  return arrayOfColumnsArrays.reduce((resultingColumns, columns) => {
-    if (columns) {
-      columns.forEach((column) => {
-        if (resultingColumns.find((resultingColumn) => resultingColumn.name === column.name)) {
-          return;
-        } else {
-          resultingColumns.push(column);
-        }
-      });
-    }
-
-    return resultingColumns;
-  }, []);
+  return uniqWith(arrayOfColumnsArrays.flat(), ({ name: a }, { name: b }) => a === b);
 }

 // This handles merging the tables produced by multiple expressions run on a single member of the `by` split.
@ -138,17 +114,17 @@ function combineAcross(datatableArray: Datatable[]) {
  });

  // Merge columns and rows.
-  const arrayOfRowsArrays = map(datatableArray, 'rows');
+  const arrayOfRowsArrays = _map(datatableArray, 'rows');
  const rows = [];
  for (let i = 0; i < targetRowLength; i++) {
-    const rowsAcross = map(arrayOfRowsArrays, i);
+    const rowsAcross = _map(arrayOfRowsArrays, i);

    // The reason for the Object.assign is that rowsAcross is an array
    // and those rows need to be applied as arguments to Object.assign
    rows.push(Object.assign({}, ...rowsAcross));
  }

-  const columns = combineColumns(map(datatableArray, 'columns'));
+  const columns = combineColumns(_map(datatableArray, 'columns'));

  return {
    type: 'datatable',
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.test.js
@ -6,6 +6,7 @@
 */

 import { of } from 'rxjs';
+import { TestScheduler } from 'rxjs/testing';
 import { functionWrapper } from '../../../test_helpers/function_wrapper';
 import { switchFn } from './switch';

@ -39,7 +40,13 @@ describe('switch', () => {
      result: 5,
    },
  ];
-  const nonMatchingCases = mockCases.filter((c) => !c.matches);
+  const nonMatchingCases = mockCases.filter(({ matches }) => !matches);
+
+  let testScheduler;
+
+  beforeEach(() => {
+    testScheduler = new TestScheduler((actual, expected) => expect(actual).toStrictEqual(expected));
+  });

  describe('spec', () => {
    it('is a function', () => {
@ -51,13 +58,19 @@ describe('switch', () => {
    describe('with no cases', () => {
      it('should return the context if no default is provided', () => {
        const context = 'foo';
-        expect(fn(context, {})).resolves.toBe(context);
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, {})).toBe('(0|)', [context])
+        );
      });

      it('should return the default if provided', () => {
        const context = 'foo';
        const args = { default: () => of('bar') };
-        expect(fn(context, args)).resolves.toBe('bar');
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', ['bar'])
+        );
      });
    });

@ -65,7 +78,10 @@ describe('switch', () => {
      it('should return the context if no default is provided', () => {
        const context = 'foo';
        const args = { case: nonMatchingCases.map(getter) };
-        expect(fn(context, args)).resolves.toBe(context);
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', [context])
+        );
      });

      it('should return the default if provided', () => {
@ -74,16 +90,22 @@ describe('switch', () => {
          case: nonMatchingCases.map(getter),
          default: () => of('bar'),
        };
-        expect(fn(context, args)).resolves.toBe('bar');
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', ['bar'])
+        );
      });
    });

    describe('with matching cases', () => {
-      it('should return the first match', async () => {
+      it('should return the first match', () => {
        const context = 'foo';
        const args = { case: mockCases.map(getter) };
-        const firstMatch = mockCases.find((c) => c.matches);
-        expect(fn(context, args)).resolves.toBe(firstMatch.result);
+        const { result } = mockCases.find(({ matches }) => matches);
+
+        testScheduler.run(({ expectObservable }) =>
+          expectObservable(fn(context, args)).toBe('(0|)', [result])
+        );
      });
    });
  });
--- a/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts
+++ b/x-pack/plugins/canvas/canvas_plugin_src/functions/common/switch.ts
@ -5,18 +5,23 @@
 * 2.0.
 */

-import { Observable } from 'rxjs';
-import { take } from 'rxjs/operators';
+import { Observable, defer, from, of } from 'rxjs';
+import { concatMap, filter, merge, pluck, take } from 'rxjs/operators';
 import { ExpressionFunctionDefinition } from 'src/plugins/expressions/common';
 import { Case } from '../../../types';
 import { getFunctionHelp } from '../../../i18n';

 interface Arguments {
-  case: Array<() => Observable<Case>>;
+  case?: Array<() => Observable<Case>>;
  default?(): Observable<any>;
 }

-export function switchFn(): ExpressionFunctionDefinition<'switch', unknown, Arguments, unknown> {
+export function switchFn(): ExpressionFunctionDefinition<
+  'switch',
+  unknown,
+  Arguments,
+  Observable<unknown>
+> {
  const { help, args: argHelp } = getFunctionHelp().switch;

  return {
@ -29,7 +34,7 @@ export function switchFn(): ExpressionFunctionDefinition<'switch', unknown, Argu
        resolve: false,
        multi: true,
        required: true,
-        help: argHelp.case,
+        help: argHelp.case!,
      },
      default: {
        aliases: ['finally'],
@ -37,18 +42,14 @@ export function switchFn(): ExpressionFunctionDefinition<'switch', unknown, Argu
        help: argHelp.default!,
      },
    },
-    fn: async (input, args) => {
-      const cases = args.case || [];
-
-      for (let i = 0; i < cases.length; i++) {
-        const { matches, result } = await cases[i]().pipe(take(1)).toPromise();
-
-        if (matches) {
-          return result;
-        }
-      }
-
-      return args.default?.().pipe(take(1)).toPromise() ?? input;
+    fn(input, args) {
+      return from(args.case ?? []).pipe(
+        concatMap((item) => item()),
+        filter(({ matches }) => matches),
+        pluck('result'),
+        merge(defer(() => args.default?.() ?? of(input))),
+        take(1)
+      );
    },
  };
 }