Replace rowGrouper with getRowGroupKey, rename groupIdGetter to getGroupId

README.md

@@ -790,7 +790,7 @@ Optional attribute to help with Cypress (or similar) selectors.
 **How it works:**
 
 1. The `groupBy` prop specifies which columns should be used for grouping
-2. The `rowGrouper` function groups rows by the specified column keys
+2. Rows are grouped by their value for each of those columns, or by `getRowGroupKey` when provided
 3. Group rows are rendered with expand/collapse toggles
 4. Child rows are nested under their parent groups
 5. Groups can be expanded/collapsed by clicking the toggle or using keyboard navigation (<kbd>←</kbd>, <kbd>→</kbd>)
@@ -840,33 +840,25 @@ const columns: readonly Column<Row>[] = [
   { key: 'name', name: 'Name' }
 ];
 
+const groupBy = ['country', 'city'];
+
 function MyGrid() {
   return (
     <TreeDataGrid
       columns={columns}
       rows={rows}
-      groupBy={['country', 'city']}
+      groupBy={groupBy}
       // ... other props
     />
   );
 }
 ```
 
-###### `rowGrouper: (rows: readonly R[], columnKey: string) => Record<string, readonly R[]>`
-
-**Required.** A function that groups rows by the specified column key. Returns an object where keys are the group values and values are arrays of rows belonging to that group.
-
-```tsx
-function rowGrouper(rows: readonly Row[], columnKey: string): Record<string, readonly Row[]> {
-  return Object.groupBy(rows, (row) => row[columnKey]);
-}
-```
-
-:warning: **Performance:** Define this function outside your component or memoize it with `useCallback` to prevent unnecessary re-renders.
+:warning: **Performance:** Define this array outside your component or memoize it with `useMemo`, otherwise all the rows are grouped again on every render.
 
 ###### `expandedGroupIds: ReadonlySet<unknown>`
 
-**Required.** A set of group IDs that are currently expanded. Group IDs are generated by `groupIdGetter`.
+**Required.** A set of group IDs that are currently expanded. Group IDs are generated by `getGroupId`.
 
 ```tsx
 import { useState } from 'react';
@@ -889,7 +881,24 @@ function MyGrid() {
 
 **Required.** Callback triggered when groups are expanded or collapsed.
 
-###### `groupIdGetter?: Maybe<(groupKey: string, parentId?: string) => string>`
+###### `getRowGroupKey?: Maybe<(row: R, columnKey: string) => string>`
+
+A function that returns the key of the group a row belongs to, for the specified column key. Rows sharing the same group key are grouped together, and groups are ordered by the first row that belongs to them. If not provided, a default implementation is used that stringifies the row's value for that column, i.e., `String(row[columnKey])`.
+
+Provide your own to group on something other than the raw cell value:
+
+```tsx
+function getRowGroupKey(row: Row, columnKey: string): string {
+  if (columnKey === 'date') {
+    return row.date.getFullYear().toString();
+  }
+  return String(row[columnKey]);
+}
+```
+
+:warning: **Performance:** Define this function outside your component or memoize it with `useCallback`, otherwise all the rows are grouped again on every render.
+
+###### `getGroupId?: Maybe<(groupKey: string, parentId?: string) => string>`
 
 Function to generate unique IDs for group rows. If not provided, a default implementation is used that concatenates parent and group keys with `__`.
 

src/TreeDataGrid.tsx

@@ -32,23 +32,42 @@ export interface TreeDataGridProps<R, SR = unknown, K extends Key = Key> extends
   columns: readonly Column<NoInfer<R>, NoInfer<SR>>[];
   rowHeight?: Maybe<number | ((args: RowHeightArgs<NoInfer<R>>) => number)>;
   groupBy: readonly string[];
-  rowGrouper: (
-    rows: readonly NoInfer<R>[],
-    columnKey: string
-  ) => Record<string, readonly NoInfer<R>[]>;
+  getRowGroupKey?: Maybe<(row: NoInfer<R>, columnKey: string) => string>;
   expandedGroupIds: ReadonlySet<unknown>;
   onExpandedGroupIdsChange: (expandedGroupIds: Set<unknown>) => void;
-  groupIdGetter?: Maybe<(groupKey: string, parentId?: string) => string>;
+  getGroupId?: Maybe<(groupKey: string, parentId?: string) => string>;
 }
 
-interface GroupBy<TRow> {
+interface PartialGroupRow<TRow> {
+  readonly groups: GroupsMap<TRow>;
+  // all the descendant rows of the group, not just the direct children
   readonly childRows: readonly TRow[];
-  readonly childGroups: readonly TRow[] | ReadonlyGroupByMap<TRow>;
   readonly startRowIndex: number;
 }
 
-type GroupByMap<TRow> = Map<string, GroupBy<TRow>>;
-type ReadonlyGroupByMap<TRow> = ReadonlyMap<string, GroupBy<TRow>>;
+type GroupsMap<TRow> = ReadonlyMap<string, PartialGroupRow<TRow>>;
+
+// only used while the rows are being grouped
+interface MutablePartialGroupRow<TRow> {
+  readonly groups: MutableGroupsMap<TRow>;
+  readonly childRows: TRow[];
+  startRowIndex: number;
+}
+
+type MutableGroupsMap<TRow> = Map<string, MutablePartialGroupRow<TRow>>;
+
+interface GroupedRows<TRow> {
+  readonly groups: GroupsMap<TRow> | undefined;
+  readonly rowsCount: number;
+}
+
+// Group rows are only ever created by `TreeDataGrid`, so every instance can share a single
+// set, which keeps `isGroupRow` stable and out of the dependency arrays below
+const allGroupRows = new WeakSet<object>();
+
+function isGroupRow<TRow>(row: TRow | GroupRow<TRow>): row is GroupRow<TRow> {
+  return typeof row === 'object' && row !== null && allGroupRows.has(row);
+}
 
 export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
   columns: rawColumns,
@@ -63,18 +82,19 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
   onSelectedRowsChange: rawOnSelectedRowsChange,
   renderers,
   groupBy: rawGroupBy,
-  rowGrouper,
+  getRowGroupKey: rawGetRowGroupKey,
   expandedGroupIds,
   onExpandedGroupIdsChange,
-  groupIdGetter: rawGroupIdGetter,
+  getGroupId: rawGetGroupId,
   ...props
 }: TreeDataGridProps<R, SR, K>) {
   const contextRenderers = useRenderers<R, SR>();
   const rawRenderRow = renderers?.renderRow ?? contextRenderers?.renderRow ?? defaultRenderRow;
   const headerAndTopSummaryRowsCount = 1 + (props.topSummaryRows?.length ?? 0);
   const { leftKey, rightKey } = getLeftRightKey(props.direction);
   const toggleGroupLatest = useLatestFunc(toggleGroup);
-  const groupIdGetter = rawGroupIdGetter ?? defaultGroupIdGetter;
+  const getRowGroupKey = rawGetRowGroupKey ?? defaultGetRowGroupKey;
+  const getGroupId = rawGetGroupId ?? defaultGetGroupId;
 
   const { columns, groupBy } = useMemo(() => {
     const columns = rawColumns.toSorted(({ key: aKey }, { key: bKey }) => {
@@ -112,59 +132,68 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     return { columns, groupBy };
   }, [rawColumns, rawGroupBy]);
 
-  const [groupedRows, rowsCount] = useMemo(() => {
-    if (groupBy.length === 0) return [undefined, rawRows.length];
+  // Grouping only depends on the rows, it is not affected by the expanded state
+  const { groups, rowsCount } = useMemo((): GroupedRows<R> => {
+    if (groupBy.length === 0) return { groups: undefined, rowsCount: rawRows.length };
+
+    // Group the rows in a single pass, each row is added to every group it belongs to
+    const groups: MutableGroupsMap<R> = new Map();
+
+    for (const row of rawRows) {
+      let parentGroups = groups;
+
+      for (const columnKey of groupBy) {
+        const groupKey = getRowGroupKey(row, columnKey);
+        const group = parentGroups.getOrInsertComputed(groupKey, () => ({
+          groups: new Map(),
+          childRows: [],
+          startRowIndex: 0
+        }));
+        group.childRows.push(row);
+        parentGroups = group.groups;
+      }
+    }
 
-    function groupRows(
-      rows: readonly R[],
-      groupByIndex: number,
-      startRowIndex: number
-    ): [ReadonlyGroupByMap<R>, number] {
+    // Number the groups as if they were all expanded, and return the total number of rows.
+    // Neither depends on the expanded state, so collapsed groups never have to be traversed
+    // when the rows are flattened.
+    function setStartRowIndexes(groups: MutableGroupsMap<R>, startRowIndex: number): number {
       let groupRowsCount = 0;
-      const groupByKey = groupBy[groupByIndex];
-      const isLastGroupBy = groupByIndex === groupBy.length - 1;
-      const groups: GroupByMap<R> = new Map();
-
-      for (const [key, childRows] of Object.entries(rowGrouper(rows, groupByKey))) {
-        // Recursively group each parent group
-        const [childGroups, childRowsCount] = isLastGroupBy
-          ? [childRows, childRows.length]
-          : groupRows(childRows, groupByIndex + 1, startRowIndex + groupRowsCount + 1); // 1 for parent row
-        groups.set(key, { childRows, childGroups, startRowIndex: startRowIndex + groupRowsCount });
-        groupRowsCount += childRowsCount + 1; // 1 for parent row
+
+      for (const group of groups.values()) {
+        group.startRowIndex = startRowIndex + groupRowsCount;
+        groupRowsCount++; // 1 for the group row
+        groupRowsCount +=
+          group.groups.size === 0
+            ? // the group is on the last level, it only contains rows
+              group.childRows.length
+            : setStartRowIndexes(group.groups, startRowIndex + groupRowsCount);
       }
 
-      return [groups, groupRowsCount];
+      return groupRowsCount;
     }
 
-    return groupRows(rawRows, 0, 0);
-  }, [groupBy, rowGrouper, rawRows]);
+    const rowsCount = setStartRowIndexes(groups, 0);
 
-  const [rows, isGroupRow] = useMemo((): [
-    readonly (R | GroupRow<R>)[],
-    (row: R | GroupRow<R>) => row is GroupRow<R>
-  ] => {
-    const allGroupRows = new Set<unknown>();
-    if (!groupedRows) return [rawRows, isGroupRow];
+    return { groups, rowsCount };
+  }, [getRowGroupKey, groupBy, rawRows]);
 
-    const flattenedRows: (R | GroupRow<R>)[] = [];
+  const rows = useMemo((): readonly (R | GroupRow<R>)[] => {
+    if (groups === undefined) {
+      return rawRows;
+    }
 
-    const expandGroup = (
-      rows: ReadonlyGroupByMap<R> | readonly R[],
-      parentId: string | undefined,
-      level: number
-    ): void => {
-      if (isReadonlyArray(rows)) {
-        flattenedRows.push(...rows);
-        return;
-      }
+    const flattenedRows: (R | GroupRow<R>)[] = [];
 
+    // Only the expanded groups are traversed
+    function expandGroups(groups: GroupsMap<R>, parentId: string | undefined, level: number): void {
       let posInSet = 0;
 
-      for (const [groupKey, row] of rows) {
-        const id = groupIdGetter(groupKey, parentId);
+      for (const [groupKey, group] of groups) {
+        // destructuring in the loop head would stop React Compiler from compiling the component
+        const { groups: childGroups, childRows, startRowIndex } = group;
+        const id = getGroupId(groupKey, parentId);
         const isExpanded = expandedGroupIds.has(id);
-        const { childRows, childGroups, startRowIndex } = row;
 
         const groupRow: GroupRow<R> = {
           id,
@@ -174,25 +203,30 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
           childRows,
           level,
           posInSet: posInSet++,
-          startRowIndex,
-          setSize: rows.size
+          setSize: groups.size,
+          startRowIndex
         };
         flattenedRows.push(groupRow);
         allGroupRows.add(groupRow);
 
         if (isExpanded) {
-          expandGroup(childGroups, id, level + 1);
+          if (childGroups.size === 0) {
+            // the group is on the last level, it only contains rows.
+            // Push them one by one: spreading them would throw in groups with too many rows
+            for (const row of childRows) {
+              flattenedRows.push(row);
+            }
+          } else {
+            expandGroups(childGroups, id, level + 1);
+          }
         }
       }
-    };
+    }
 
-    expandGroup(groupedRows, undefined, 0);
-    return [flattenedRows, isGroupRow];
+    expandGroups(groups, undefined, 0);
 
-    function isGroupRow(row: R | GroupRow<R>): row is GroupRow<R> {
-      return allGroupRows.has(row);
-    }
-  }, [expandedGroupIds, groupedRows, rawRows, groupIdGetter]);
+    return flattenedRows;
+  }, [expandedGroupIds, getGroupId, groups, rawRows]);
 
   const rowHeight = useMemo(() => {
     if (typeof rawRowHeight === 'function') {
@@ -205,7 +239,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     }
 
     return rawRowHeight;
-  }, [isGroupRow, rawRowHeight]);
+  }, [rawRowHeight]);
 
   const getParentRowAndIndex = useCallback(
     (row: R | GroupRow<R>) => {
@@ -219,7 +253,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
       return undefined;
     },
-    [isGroupRow, rows]
+    [rows]
   );
 
   const rowKeyGetter = useCallback(
@@ -241,7 +275,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
       return rows.indexOf(row);
     },
-    [getParentRowAndIndex, isGroupRow, rawRowKeyGetter, rows]
+    [getParentRowAndIndex, rawRowKeyGetter, rows]
   );
 
   const selectedRows = useMemo((): Maybe<ReadonlySet<Key>> => {
@@ -263,7 +297,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     }
 
     return selectedRows;
-  }, [isGroupRow, rawRowKeyGetter, rawSelectedRows, rows]);
+  }, [rawRowKeyGetter, rawSelectedRows, rows]);
 
   function onSelectedRowsChange(newSelectedRows: Set<Key>) {
     if (!rawOnSelectedRowsChange) return;
@@ -453,10 +487,10 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
   );
 }
 
-function defaultGroupIdGetter(groupKey: string, parentId: string | undefined) {
-  return parentId === undefined ? groupKey : `${parentId}__${groupKey}`;
+function defaultGetRowGroupKey(row: unknown, columnKey: string) {
+  return String((row as Record<string, unknown>)[columnKey]);
 }
 
-function isReadonlyArray(arr: unknown): arr is readonly unknown[] {
-  return Array.isArray(arr);
+function defaultGetGroupId(groupKey: string, parentId: string | undefined) {
+  return parentId === undefined ? groupKey : `${parentId}__${groupKey}`;
 }