Nested grouping by one or more levels-CodePudding

Given the following data structure:

const data = [
  { A: 1, B: 12, C: 123 },
  { A: 1, B: 122, C: 1233 },
  { A: 2, B: 22, C: 223 }
];

I would like to implement a function called groupBy with the following signature:

function groupBy<T, By extends keyof T>(object: T[], by: ...By[]): Map<By[0], Map<By[1], ..., Map<number, T[]>...>

which I would be able to call like:

const A = groupBy(data, "A"); // Map<number, { A: number, B: number, C: number }[]>;
const AB = groupBy(data, "A", "B"); // Map<number, Map<number, { A: number, B: number, C: number }[]>>;
const ABC = groupBy(data, "A", "B", "C"); // Map<number, Map<number, Map<number, ...>>;

Unfortunately, I was only able to implement groupBy with a single level:

function groupBy<T extends object, K extends keyof T>(collection: T[], iteratee: K): Map<T[K], T[]> {
  const map: Map<T[K], T[]> = new Map();

  for (const item of collection) {
    const accumalated = map.get(item[iteratee]);
    if (accumalated === undefined) {
      map.set(item[iteratee], [item]);
    } else {
      map.set(item[iteratee], [...accumalated, item]);
    }
  }

  return map;
}

CodePudding user response：

First let's describe groupBy() at the type level, by giving it a strongly-typed call signature:

type GroupedBy<T, K> = K extends [infer K0, ...infer KR] ?
    Map<T[Extract<K0, keyof T>], GroupedBy<T, KR>> : T[];

// call signature
function groupBy<T, K extends Array<keyof T>>(
  objects: readonly T[], ...by: [...K]
): GroupedBy<T, K>;

So groupBy() is generic in T, the type of the elements of the objects array, as well as K, a tuple of the keys of T corresponding to the ...by rest parameter. The function returns GroupedBy<T, K>.

So what's GroupedBy<T, K>? It's a recursive conditional type. If the tuple K is empty, then it will just be T[] (since grouping an array by nothing should yield the same array). Otherwise, we use variadic tuple types to split the K tuple into the first element, K0, and the rest, KR. Then GroupedBy<T, K> will be a Map whose key type is the type of property of T at key K0 (conceptually this is just an indexed access type T[K0], but the compiler doesn't know that K0 will be a key of T, so we use the Extract<T, U> utility type to convince it... so T[Extract<K0, keyof T>]) and whose value type is, recursing, GroupedBy<T, KR>.

Let's make sure the compiler does the right thing:

const A = groupBy(data, "A"); 
// Map<number, { A: number, B: number, C: number }[]>;

const AB = groupBy(data, "A", "B"); 
// Map<number, Map<number, { A: number, B: number, C: number }[]>>;

const ABC = groupBy(data, "A", "B", "C"); /* Map<number, Map<number, Map<number, {
    A: number;
    B: number;
    C: number;
}[]>>> */

Looks good. Just to be sure, let's change number into something else:

const otherData = [
    { str: "a", num: 1, bool: true },
    { str: "a", num: 1, bool: false },
    { str: "a", num: 2, bool: true }
];
const grouped = groupBy(otherData, "str", "num", "bool")
/* const grouped: Map<string, Map<number, Map<boolean, {
    str: string;
    num: number;
    bool: boolean;
}[]>>> */

Also looks good.

Now let's implement groupBy(). The compiler can't possibly follow the recursive conditional typing inside the implementation (see microsoft/TypeScript#33912), so let's loosen things up by making it an overloaded function with a single, strongly-typed call signature, and a loose implementation signature using the any type. We have to be careful to do it right because the compiler won't catch type bugs.

Anyway, here's a possible implementation:

// implementation
function groupBy(objects: readonly any[], ...by: Array<PropertyKey>) {
    if (!by.length) return objects;
    const [k0, ...kr] = by;
    const topLevelGroups = new Map<any, any[]>();
    for (const obj of objects) {
        let k = obj[k0];
        let arr = topLevelGroups.get(k);
        if (!arr) {
            arr = [];
            topLevelGroups.set(k, arr);
        }
        arr.push(obj);
    }
    return new Map(Array.from(topLevelGroups, ([k, v]) => ([k, groupBy(v, ...kr)])));

}

If the by array is empty, return the objects array unchanged; this corresponds to the base case of GroupedBy<T, K> where K is []. Otherwise, we split by into the first element k0 and the rest of it kr. We then go through and do a top-level grouping of objects by the values in their k0 key. This involves getting the right making sure that we initialize arrays before pushing things into them. And finally, at the end, we transform that top-level grouping (using this technique), recursing, applying groupBy() to each array of objects.

Let's see if it works:

console.log(A);
/* Map (2) {1 => [{
  "A": 1,
  "B": 12,
  "C": 123
}, {
  "A": 1,
  "B": 122,
  "C": 1233
}], 2 => [{
  "A": 2,
  "B": 22,
  "C": 223
}]}  */

console.log(AB);
/* Map (2) {1 => Map (2) {12 => [{
  "A": 1,
  "B": 12,
  "C": 123
}], 122 => [{
  "A": 1,
  "B": 122,
  "C": 1233
}]}, 2 => Map (1) {22 => [{
  "A": 2,
  "B": 22,
  "C": 223
}]}}  */

Also looks good.

Playground link to code