Introducing switchPartition: A Versatile RxJS Operator for Stream Partitioning

[omarzougui]

The switchPartition operator aims to extend the capabilities of RxJS by providing a more flexible and type-safe way to partition an observable stream into multiple observables based on specified cases. This operator is particularly useful in scenarios where you need to segregate a stream of events or data into different categories for further processing or side-effects.

For example, consider a stream of vehicles that could be Electric, Hybrid, Gasoline, or Diesel. Using switchPartition, you can easily create four separate streams for each type of vehicle. This makes it easier to handle different logic for each category without having to filter the same original stream multiple times.

const [
    electric$, 
    hybrid$,
    gasoline$,
    diesel$
] = switchPartition(cars$, ['Electric', 'Hybrid', 'Gasoline', 'Diesel']);

The operator also supports an optional projection function, allowing for more complex use-cases, such as partitioning based on a property of the emitted values.

const [feature$, bugfix$, docs$, chore$] = switchPartition(
    pullRequest$, 
    ['Feature', 'Bugfix', 'Docs', 'Chore'],
    (pr: PullRequest) => pr.type
);

The switchPartition operator is designed to be highly type-safe, ensuring that the types of the emitted values in the partitioned observables match the types specified in the cases array. This makes it a robust and valuable addition to the RxJS library.

import { filter, Observable } from "rxjs";

type Projection<T, P> = (value: T) => P;
type _ArrayOfSize<
  T,
  N extends number,
  R extends unknown[],
> = R["length"] extends N ? R : _ArrayOfSize<T, N, [T, ...R]>;
export type StreamCases<T, N extends number> = N extends N
  ? number extends N
    ? T[]
    : _ArrayOfSize<T, N, []>
  : never;
type TupleTag = { __tuple__: true };
type IsTuple<T> = T extends readonly [any, ...infer Rest]
  ? true
  : T extends TupleTag
  ? true
  : false;
type EnsureTuple<T> = T & (IsTuple<T> extends true ? unknown : never);

export function switchPartition<T, K1 extends T, K2 extends T>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2]>>,
  thisArg?: any,
): [Observable<K1>, Observable<K2>];
export function switchPartition<T, K1, K2, P extends K1 | K2>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [Observable<T>, Observable<T>];
export function switchPartition<T, K1 extends T, K2 extends T, K3 extends T>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3]>>,
  thisArg?: any,
): [Observable<K1>, Observable<K2>, Observable<K3>];
export function switchPartition<T, K1, K2, K3, P extends K1 | K2 | K3>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [Observable<T>, Observable<T>, Observable<T>];
export function switchPartition<
  T,
  K1 extends T,
  K2 extends T,
  K3 extends T,
  K4 extends T,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4]>>,
  thisArg?: any,
): [Observable<K1>, Observable<K2>, Observable<K3>, Observable<K4>];
export function switchPartition<T, K1, K2, K3, K4, P extends K1 | K2 | K3 | K4>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [Observable<T>, Observable<T>, Observable<T>, Observable<T>];
export function switchPartition<
  T,
  K1 extends T,
  K2 extends T,
  K3 extends T,
  K4 extends T,
  K5 extends T,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5]>>,
  thisArg?: any,
): [
  Observable<K1>,
  Observable<K2>,
  Observable<K3>,
  Observable<K4>,
  Observable<K5>,
];
export function switchPartition<
  T,
  K1,
  K2,
  K3,
  K4,
  K5,
  P extends K1 | K2 | K3 | K4 | K5,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [Observable<T>, Observable<T>, Observable<T>, Observable<T>, Observable<T>];
export function switchPartition<
  T,
  K1 extends T,
  K2 extends T,
  K3 extends T,
  K4 extends T,
  K5 extends T,
  K6 extends T,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5, K6]>>,
  thisArg?: any,
): [
  Observable<K1>,
  Observable<K2>,
  Observable<K3>,
  Observable<K4>,
  Observable<K5>,
  Observable<K6>,
];
export function switchPartition<
  T,
  K1,
  K2,
  K3,
  K4,
  K5,
  K6,
  P extends K1 | K2 | K3 | K4 | K5 | K6,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5, K6]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
];
export function switchPartition<
  T,
  K1 extends T,
  K2 extends T,
  K3 extends T,
  K4 extends T,
  K5 extends T,
  K6 extends T,
  K7 extends T,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5, K6, K7]>>,
  thisArg?: any,
): [
  Observable<K1>,
  Observable<K2>,
  Observable<K3>,
  Observable<K4>,
  Observable<K5>,
  Observable<K6>,
  Observable<K7>,
];
export function switchPartition<
  T,
  K1,
  K2,
  K3,
  K4,
  K5,
  K6,
  K7,
  P extends K1 | K2 | K3 | K4 | K5 | K6 | K7,
>(
  source$: Observable<T>,
  cases: Readonly<EnsureTuple<[K1, K2, K3, K4, K5, K6, K7]>>,
  projection: Projection<T, P>,
  thisArg?: any,
): [
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
  Observable<T>,
];
/**
 * General implementation of the `switchPartition` function.
 *
 * @example
 * const [electric$, hybrid$, gasoline$, diesel$] = switchPartition(cars$, ['Electric', 'Hybrid', 'Gasoline', 'Diesel']);
 * const [feature$, bugfix$, docs$, chore$] = switchPartition(pullRequests, ['Feature', 'Bugfix', 'Docs', 'Chore'], (pr: PullRequest) => pr.type)
 * @param source$ - The observable source stream.
 * @param cases - An array of cases to filter the stream.
 * @param projection - Optional projection function to project the source type to a different type.
 * @returns An array of observables, one for each case.
 * @param thisArg An optional argument to determine the value of `this` in the `projection` function.
 */
export function switchPartition<T, K extends any[], P extends K[number]>(
  source$: Observable<T>,
  cases: K,
  projection?: Projection<T, P>,
  thisArg?: any,
): Observable<T>[] {
  return cases.map((specificCase) =>
    source$.pipe(
      filter(
        (value) => specificCase === (projection?.call(thisArg, value) ?? value),
      ),
    ),
  );
}

[voliva]

It's an interesting operator and I can see the effort, but I think it has some shortcomings. Please, don't take this as destructive critsism 🙏

1. The types feel a bit overkill. An Observable of a primitive literal (e.g. banana$: Observable<'Banana'>) carries the same amount of information than an Observable of void (the example would become banana$: Observable<void>). Types would be great if it can discriminate against non-primitives: Partition Observable<SuccessResponse | ErrorResponse> into two observables Observable<SuccessResponse> and Observable<ErrorResponse>.
2. The source$ is not being shared, which means that each subscriber to each individual "sub-observable" will trigger a new subscription to the source.
3. When source$ emits, it will run the filter function for each one of the observables. This is different from a "partition" or a "groupBy" operator, where it doesn't need to run the filter multiple times.

In React-RxJS we have done one operator that deals with this, published into a separate package @react-rxjs/utils https://react-rxjs.org/docs/api/utils/partitionByKey (source, tests)

It doesn't discriminate union types, but it has semantics defined on when a group becomes alive and when it disappears. The way this works is that you access the group through a function. Your example on PRs would become:

const [prByType$] = partitionByKey(
  pullRequest$,
  pr => pr.type;
);

//  Later, when consuming
prByType$('Feature').pipe(...)

It returns that function within a tuple, because it also returns an additional observable that emits which keys become active or which ones get removed, for cases where you don't know the keys beforehand (e.g. you want to have a partition by User ID)

[omarzougui]

Thank you for taking the time to thoroughly read, analyze, and review my code.

We agree that the partition b ke "partitionByKKe()" partitions the streammm only into 2 and not into n streams ?
Excessive typing is to avoid the number of derived streams being different from the number of boxes
basically what bothers me today with the partition is that if I want to partition my source stream into 6 different streams I have to use the partition or filter operator 5 or 6 times.

For typing it can become less restrictive if we add another layer of overload to accept an array of cases without notion of length
But I find telling the user that you cannot create 6 observables from a source observable that you have partitioned with only 5 cases useful

export function switchPartition<T, K extends any[], P extends K[number]>(
  source$: Observable<T>,
  cases: K,
  projection?: Projection<T, P>,
  thisArg?: any,
): Observable<T>[] {
  return cases.map((specificCase) =>
    source$.pipe(
      share(),
      filter(
        (value) => specificCase === (projection?.call(thisArg, value) ?? value),
      ),
    ),
  );
}

[voliva]

Ah, I see your use case.

Excessive typing is to avoid the number of derived streams being different from the number of boxes
basically what bothers me today with the partition is that if I want to partition my source stream into 6 different streams I have to use the partition or filter operator 5 or 6 times.

I would also suggest

type MappedObservables<InputKeys extends readonly any[], Result> = {
  [K in keyof InputKeys]: Observable<Result>
};

export function switchPartition<T, K extends readonly any[], P extends K[number]>(
  source$: Observable<T>,
  cases: K,
  projection?: Projection<T, P>,
  thisArg?: any,
): MappedObservables<K, T> {
  return cases.map((specificCase) =>
    source$.pipe(
      filter(
        (value) => specificCase === (projection?.call(thisArg, value) ?? value),
      ),
    ),
  );
}

It has some tradeoffs though, TS has a harder time treating each value as individual literals

Be careful on the usage of share() though, I actually meant in the source$ itself. In a way, I wouldn't expect the following to trigger two subscriptions on the source$:

const [
    electric$, 
    hybrid$,
    gasoline$,
    diesel$
] = switchPartition(cars$, ['Electric', 'Hybrid', 'Gasoline', 'Diesel']);

const ecological$ = merge(electric$, hybrid$);
ecological$.subscribe(....)

If you add share() within each mapped case, then you're only sharing the subscription within that specific group. Which means that the code above will still trigger two subscriptions to cars$

We agree that the partition b ke "partitionByKKe()" partitions the streammm only into 2 and not into n streams ?

No, partitionByKey handles as many streams as you want. And it can handle also as many as you want without knowing which ones you will have beforehand (such as the example I said of using a userId as the key). This isn't something you can do with the current switchPartition.

partitionByKey the way it solves the shared stream / late subscribers is that it will keep each subgroup "alive", storing the latest value for each one of them. So if someone subscribes late, they will receive the last value for that specific group.

But it means that they are stateful, so it's not suitable for events. Maybe you can just move the share() before doing each map if that's what you need.

[Laurens-makel]

I think my feature request is very similair #7420.

Mine started out with the idea of partition on base of index returned but later added support for a string array as well.