fix: don't crash the table on 1 million rows

frontend/.storybook/preview.tsx

@@ -7,7 +7,7 @@ import React, { useEffect } from "react";
 import { cn } from "../src/utils/cn";
 import { TooltipProvider } from "../src/components/ui/tooltip";
 import { Toaster } from "../src/components/ui/toaster";
-import { TailwindIndicator } from "../src/components/ui/tailwind-indicator";
+import { TailwindIndicator } from "../src/components/indicator";
 
 const withTheme: Decorator = (Story, context) => {
   const theme = context.globals.theme || "light";

frontend/src/components/data-table/__test__/columns.test.ts

@@ -0,0 +1,23 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { expect, test } from "vitest";
+import { uniformSample } from "../columns";
+
+test("uniformSample", () => {
+  const items = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"];
+
+  expect(uniformSample(items, 2)).toMatchInlineSnapshot(`
+    [
+      "A",
+      "J",
+    ]
+  `);
+  expect(uniformSample(items, 4)).toMatchInlineSnapshot(`
+    [
+      "A",
+      "C",
+      "F",
+      "J",
+    ]
+  `);
+  expect(uniformSample(items, 100)).toBe(items);
+});

frontend/src/components/data-table/columns.tsx

@@ -9,7 +9,7 @@ interface ColumnInfo {
   type: "primitive" | "mime";
 }
 
-export function getColumnInfo<T>(items: T[]): ColumnInfo[] {
+function getColumnInfo<T>(items: T[]): ColumnInfo[] {
   // No items
   if (items.length === 0) {
     return [];
@@ -21,7 +21,10 @@ export function getColumnInfo<T>(items: T[]): ColumnInfo[] {
   }
 
   const keys = new Map<string, ColumnInfo>();
-  items.forEach((item) => {
+
+  // This can be slow for large datasets,
+  // so only sample 10 evenly distributed rows
+  uniformSample(items, 10).forEach((item) => {
     if (typeof item !== "object") {
       return;
     }
@@ -153,3 +156,21 @@ function isPrimitiveOrNullish(value: unknown): boolean {
   const isObject = typeof value === "object";
   return !isObject;
 }
+
+/**
+ * Uniformly sample n items from an array
+ */
+export function uniformSample<T>(items: T[], n: number): T[] {
+  if (items.length <= n) {
+    return items;
+  }
+  const sample: T[] = [];
+  const step = items.length / n;
+  for (let i = 0; i < n - 1; i++) {
+    const idx = Math.floor(i * step);
+    sample.push(items[idx]);
+  }
+  const last = items.at(-1);
+  sample.push(last!);
+  return sample;
+}

frontend/src/components/data-table/data-table.tsx

@@ -1,5 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import React from "react";
+import React, { memo } from "react";
 import {
   ColumnDef,
   OnChangeFn,
@@ -37,99 +37,101 @@ interface DataTableProps<TData, TValue> extends Partial<DownloadActionProps> {
   onRowSelectionChange?: OnChangeFn<RowSelectionState>;
 }
 
-export const DataTable = <TData, TValue>({
-  wrapperClassName,
-  className,
-  columns,
-  data,
-  rowSelection,
-  pageSize = 10,
-  downloadAs,
-  pagination = false,
-  onRowSelectionChange,
-}: DataTableProps<TData, TValue>) => {
-  const [sorting, setSorting] = React.useState<SortingState>([]);
-  const [paginationState, setPaginationState] = React.useState<PaginationState>(
-    { pageSize: pageSize, pageIndex: 0 },
-  );
-
-  const table = useReactTable({
-    data,
+export const DataTable = memo(
+  <TData, TValue>({
+    wrapperClassName,
+    className,
     columns,
-    getCoreRowModel: getCoreRowModel(),
-    // pagination
-    onPaginationChange: setPaginationState,
-    getPaginationRowModel: pagination ? getPaginationRowModel() : undefined,
-    // sorting
-    onSortingChange: setSorting,
-    getSortedRowModel: getSortedRowModel(),
-    // selection
-    onRowSelectionChange: onRowSelectionChange,
-    state: {
-      sorting,
-      pagination: pagination
-        ? { ...paginationState, pageSize: pageSize }
-        : { pageIndex: 0, pageSize: data.length },
-      rowSelection,
-    },
-  });
+    data,
+    rowSelection,
+    pageSize = 10,
+    downloadAs,
+    pagination = false,
+    onRowSelectionChange,
+  }: DataTableProps<TData, TValue>) => {
+    const [sorting, setSorting] = React.useState<SortingState>([]);
+    const [paginationState, setPaginationState] =
+      React.useState<PaginationState>({ pageSize: pageSize, pageIndex: 0 });
 
-  return (
-    <div className={cn(wrapperClassName, "flex flex-col space-y-2")}>
-      <div className={cn(className || "rounded-md border")}>
-        <Table>
-          <TableHeader>
-            {table.getHeaderGroups().map((headerGroup) => (
-              <TableRow key={headerGroup.id}>
-                {headerGroup.headers.map((header) => {
-                  return (
-                    <TableHead key={header.id}>
-                      {header.isPlaceholder
-                        ? null
-                        : flexRender(
-                            header.column.columnDef.header,
-                            header.getContext(),
-                          )}
-                    </TableHead>
-                  );
-                })}
-              </TableRow>
-            ))}
-          </TableHeader>
-          <TableBody>
-            {table.getRowModel().rows?.length ? (
-              table.getRowModel().rows.map((row) => (
-                <TableRow
-                  key={row.id}
-                  data-state={row.getIsSelected() && "selected"}
-                >
-                  {row.getVisibleCells().map((cell) => (
-                    <TableCell key={cell.id}>
-                      {flexRender(
-                        cell.column.columnDef.cell,
-                        cell.getContext(),
-                      )}
-                    </TableCell>
-                  ))}
+    const table = useReactTable({
+      data,
+      columns,
+      getCoreRowModel: getCoreRowModel(),
+      // pagination
+      onPaginationChange: setPaginationState,
+      getPaginationRowModel: pagination ? getPaginationRowModel() : undefined,
+      // sorting
+      onSortingChange: setSorting,
+      getSortedRowModel: getSortedRowModel(),
+      // selection
+      onRowSelectionChange: onRowSelectionChange,
+      state: {
+        sorting,
+        pagination: pagination
+          ? { ...paginationState, pageSize: pageSize }
+          : { pageIndex: 0, pageSize: data.length },
+        rowSelection,
+      },
+    });
+
+    return (
+      <div className={cn(wrapperClassName, "flex flex-col space-y-2")}>
+        <div className={cn(className || "rounded-md border")}>
+          <Table>
+            <TableHeader>
+              {table.getHeaderGroups().map((headerGroup) => (
+                <TableRow key={headerGroup.id}>
+                  {headerGroup.headers.map((header) => {
+                    return (
+                      <TableHead key={header.id}>
+                        {header.isPlaceholder
+                          ? null
+                          : flexRender(
+                              header.column.columnDef.header,
+                              header.getContext(),
+                            )}
+                      </TableHead>
+                    );
+                  })}
                 </TableRow>
-              ))
-            ) : (
-              <TableRow>
-                <TableCell
-                  colSpan={columns.length}
-                  className="h-24 text-center"
-                >
-                  No results.
-                </TableCell>
-              </TableRow>
-            )}
-          </TableBody>
-        </Table>
-      </div>
-      <div className="flex align-items justify-between flex-shrink-0">
-        {pagination ? <DataTablePagination table={table} /> : <div />}
-        {downloadAs && <DownloadAs downloadAs={downloadAs} />}
+              ))}
+            </TableHeader>
+            <TableBody>
+              {table.getRowModel().rows?.length ? (
+                table.getRowModel().rows.map((row) => (
+                  <TableRow
+                    key={row.id}
+                    data-state={row.getIsSelected() && "selected"}
+                  >
+                    {row.getVisibleCells().map((cell) => (
+                      <TableCell key={cell.id}>
+                        {flexRender(
+                          cell.column.columnDef.cell,
+                          cell.getContext(),
+                        )}
+                      </TableCell>
+                    ))}
+                  </TableRow>
+                ))
+              ) : (
+                <TableRow>
+                  <TableCell
+                    colSpan={columns.length}
+                    className="h-24 text-center"
+                  >
+                    No results.
+                  </TableCell>
+                </TableRow>
+              )}
+            </TableBody>
+          </Table>
+        </div>
+        <div className="flex align-items justify-between flex-shrink-0">
+          {pagination ? <DataTablePagination table={table} /> : <div />}
+          {downloadAs && <DownloadAs downloadAs={downloadAs} />}
+        </div>
       </div>
-    </div>
-  );
-};
+    );
+  },
+);
+DataTable.displayName = "DataTable";

frontend/src/components/data-table/pagination.tsx

@@ -25,14 +25,14 @@ export const DataTablePagination = <TData,>({
     if (isAllPageSelected && !isAllSelected) {
       return (
         <span>
-          {selected} selected
+          {prettyNumber(selected)} selected
           <Button
             size="xs"
             data-testid="select-all-button"
             variant="link"
             onClick={() => table.toggleAllRowsSelected(true)}
           >
-            Select all {count}
+            Select all {prettyNumber(count)}
           </Button>
         </span>
       );
@@ -41,7 +41,7 @@ export const DataTablePagination = <TData,>({
     if (selected) {
       return (
         <span>
-          {selected} selected
+          {prettyNumber(selected)} selected
           <Button
             size="xs"
             data-testid="clear-selection-button"
@@ -54,7 +54,7 @@ export const DataTablePagination = <TData,>({
       );
     }
 
-    return `${count} items`;
+    return `${prettyNumber(count)} items`;
   };
   const currentPage = Math.min(
     table.getState().pagination.pageIndex + 1,