@marimo-team/frontend 0.21.2-dev77 → 0.21.2-dev79

package/src/components/data-table/renderers.tsx

@@ -11,7 +11,8 @@ import {
   type Table,
   type Table as TanStackTable,
 } from "@tanstack/react-table";
-import { type JSX, useRef } from "react";
+import { useVirtualizer } from "@tanstack/react-virtual";
+import { type JSX, useLayoutEffect, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import {
   TableBody,
@@ -27,6 +28,7 @@ import { DataTableContextMenu } from "./context-menu";
 import { CellRangeSelectionIndicator } from "./range-focus/cell-selection-indicator";
 import { useCellRangeSelection } from "./range-focus/use-cell-range-selection";
 import { useScrollIntoViewOnFocus } from "./range-focus/use-scroll-into-view";
+import { TABLE_ROW_HEIGHT_PX } from "./types";
 import { stringifyUnknownValue } from "./utils";
 
 export function renderTableHeader<TData>(
@@ -79,6 +81,7 @@ interface DataTableBodyProps<TData> {
   rowViewerPanelOpen: boolean;
   getRowIndex?: (row: TData, idx: number) => number;
   viewedRowIdx?: number;
+  virtualize?: boolean;
 }
 
 export const DataTableBody = <TData,>({
@@ -87,9 +90,30 @@ export const DataTableBody = <TData,>({
   rowViewerPanelOpen,
   getRowIndex,
   viewedRowIdx,
+  virtualize = false,
 }: DataTableBodyProps<TData>) => {
-  // Automatically scroll focused cells into view
+  const rows = table.getRowModel().rows;
+
+  // Find the scroll container (tbody -> table -> overflow-auto wrapper div).
+  // Using useState so that when the element becomes available after mount,
+  // useVirtualizer re-observes the correct element.
+  const [scrollElement, setScrollElement] = useState<HTMLElement | null>(null);
   const tableRef = useRef<HTMLTableSectionElement>(null);
+  useLayoutEffect(() => {
+    // tbody.parentElement = table, table.parentElement = overflow wrapper
+    setScrollElement(tableRef.current?.parentElement?.parentElement ?? null);
+  }, []);
+
+  // Always call useVirtualizer (rules of hooks); count=0 when not virtualizing
+  const virtualizer = useVirtualizer({
+    count: virtualize ? rows.length : 0,
+    getScrollElement: () => scrollElement,
+    estimateSize: () => TABLE_ROW_HEIGHT_PX,
+    overscan: 10,
+  });
+
+  // Automatically scroll focused cells into view.
+  // In virtual mode, off-screen cells won't be in the DOM so this silently no-ops for them.
   useScrollIntoViewOnFocus(tableRef);
 
   const {
@@ -172,57 +196,90 @@ export const DataTableBody = <TData,>({
 
   const hoverTemplate = table.getState().cellHoverTemplate || null;
 
-  const tableBody = (
-    <TableBody onKeyDown={handleCellsKeyDown} ref={tableRef}>
-      {table.getRowModel().rows?.length ? (
-        table.getRowModel().rows.map((row) => {
-          // Only find the row index if the row viewer panel is open
-          const rowIndex = rowViewerPanelOpen
-            ? (getRowIndex?.(row.original, row.index) ?? row.index)
-            : undefined;
-          const isRowViewedInPanel =
-            rowViewerPanelOpen && viewedRowIdx === rowIndex;
+  const renderRow = (row: Row<TData>) => {
+    // Only find the row index if the row viewer panel is open
+    const rowIndex = rowViewerPanelOpen
+      ? (getRowIndex?.(row.original, row.index) ?? row.index)
+      : undefined;
+    const isRowViewedInPanel = rowViewerPanelOpen && viewedRowIdx === rowIndex;
+
+    // Compute hover title once per row using all visible cells
+    let rowTitle: string | undefined;
+    if (hoverTemplate) {
+      const visibleCells = row.getVisibleCells?.() ?? [
+        ...row.getLeftVisibleCells(),
+        ...row.getCenterVisibleCells(),
+        ...row.getRightVisibleCells(),
+      ];
+      rowTitle = applyHoverTemplate(hoverTemplate, visibleCells);
+    }
 
-          // Compute hover title once per row using all visible cells
-          let rowTitle: string | undefined;
-          if (hoverTemplate) {
-            const visibleCells = row.getVisibleCells?.() ?? [
-              ...row.getLeftVisibleCells(),
-              ...row.getCenterVisibleCells(),
-              ...row.getRightVisibleCells(),
-            ];
-            rowTitle = hoverTemplate
-              ? applyHoverTemplate(hoverTemplate, visibleCells)
-              : undefined;
-          }
+    return (
+      <TableRow
+        key={row.id}
+        data-state={row.getIsSelected() && "selected"}
+        title={rowTitle}
+        // These classes ensure that empty rows (nulls) still render
+        className={cn(
+          "border-t h-6",
+          rowViewerPanelOpen && "cursor-pointer",
+          isRowViewedInPanel &&
+            "bg-(--blue-3) hover:bg-(--blue-3) data-[state=selected]:bg-(--blue-4)",
+        )}
+        onClick={() => handleRowClick(row)}
+      >
+        {renderCells(row.getLeftVisibleCells())}
+        {renderCells(row.getCenterVisibleCells())}
+        {renderCells(row.getRightVisibleCells())}
+      </TableRow>
+    );
+  };
 
-          return (
-            <TableRow
-              key={row.id}
-              data-state={row.getIsSelected() && "selected"}
-              title={rowTitle}
-              // These classes ensure that empty rows (nulls) still render
-              className={cn(
-                "border-t h-6",
-                rowViewerPanelOpen && "cursor-pointer",
-                isRowViewedInPanel &&
-                  "bg-(--blue-3) hover:bg-(--blue-3) data-[state=selected]:bg-(--blue-4)",
-              )}
-              onClick={() => handleRowClick(row)}
-            >
-              {renderCells(row.getLeftVisibleCells())}
-              {renderCells(row.getCenterVisibleCells())}
-              {renderCells(row.getRightVisibleCells())}
-            </TableRow>
-          );
-        })
-      ) : (
+  const renderRows = () => {
+    if (rows.length === 0) {
+      return (
         <TableRow>
           <TableCell colSpan={columns.length} className="h-24 text-center">
             No results.
           </TableCell>
         </TableRow>
-      )}
+      );
+    }
+
+    if (virtualize) {
+      const virtualItems = virtualizer.getVirtualItems();
+      const totalSize = virtualizer.getTotalSize();
+      return (
+        <>
+          {virtualItems[0]?.start > 0 && (
+            <tr
+              data-virtual-spacer=""
+              style={{ height: virtualItems[0].start }}
+            >
+              <td colSpan={columns.length} />
+            </tr>
+          )}
+          {virtualItems.map((vItem) => renderRow(rows[vItem.index]))}
+          {virtualItems.length > 0 && (
+            <tr
+              data-virtual-spacer=""
+              style={{
+                height: totalSize - (virtualItems.at(-1)?.end ?? totalSize),
+              }}
+            >
+              <td colSpan={columns.length} />
+            </tr>
+          )}
+        </>
+      );
+    }
+
+    return rows.map((row) => renderRow(row));
+  };
+
+  const tableBody = (
+    <TableBody onKeyDown={handleCellsKeyDown} ref={tableRef}>
+      {renderRows()}
     </TableBody>
   );
 

package/src/components/data-table/types.ts

@@ -10,6 +10,20 @@ declare module "@tanstack/react-table" {
   }
 }
 
+// Pixel heights derived from Tailwind classes applied to table elements.
+// row: h-6 = 24px (TableRow in renderers.tsx)
+// header: min-h-10 = 40px (TableHead in renderers.tsx)
+export const TABLE_ROW_HEIGHT_PX = 24;
+export const TABLE_HEADER_HEIGHT_PX = 40;
+
+// Default number of visible rows when virtualizing without an explicit maxHeight.
+export const DEFAULT_VIRTUAL_ROWS = 15;
+
+// Minimum row count before virtualization kicks in. Below this threshold the
+// DOM overhead is negligible and the virtualizer's measurement cost isn't
+// worth it. Must be greater than DEFAULT_VIRTUAL_ROWS.
+export const MIN_ROWS_TO_VIRTUALIZE = 100;
+
 export type ColumnName = string;
 
 export const ColumnHeaderStatsKeys = [

package/src/core/islands/main.ts

@@ -113,7 +113,6 @@ export async function initialize() {
       case "installing-package-alert":
       case "completion-result":
       case "reload":
-      case "update-cell-codes":
       case "update-cell-ids":
       case "focus-cell":
       case "variables":

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -97,7 +97,7 @@ export function useMarimoKernelConnection(opts: {
   const { autoInstantiate, sessionId, setCells } = opts;
   const { showBoundary } = useErrorBoundary();
 
-  const { handleCellMessage, setCellCodes, setCellIds } = useCellActions();
+  const { handleCellMessage, setCellIds } = useCellActions();
   const actionsWithoutMiddleware = useCellActions({ skipMiddleware: true });
 
   const handleDocumentTransaction = (
@@ -328,15 +328,6 @@ export function useMarimoKernelConnection(opts: {
       case "focus-cell":
         focusAndScrollCellOutputIntoView(msg.data.cell_id);
         return;
-      case "update-cell-codes":
-        setCellCodes({
-          codes: msg.data.codes,
-          ids: msg.data.cell_ids,
-          codeIsStale: msg.data.code_is_stale,
-          names: msg.data.names,
-          configs: msg.data.configs,
-        });
-        return;
       case "update-cell-ids":
         setCellIds({ cellIds: msg.data.cell_ids });
         return;