Add multi-column sorting with stack-based behavior

.claude/settings.local.json

@@ -0,0 +1,3 @@
+{
+  "outputStyle": "friendly-educational"
+}

frontend/src/components/data-table/__tests__/header-items.test.tsx

@@ -0,0 +1,117 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { SortingState } from "@tanstack/react-table";
+import { describe, expect, it, vi } from "vitest";
+
+describe("multi-column sorting logic", () => {
+  // Extract the core sorting logic to test in isolation
+  const handleSort = (
+    columnId: string,
+    desc: boolean,
+    sortingState: SortingState,
+    setSorting: (state: SortingState) => void,
+    clearSorting: () => void,
+  ) => {
+    const currentSort = sortingState.find((s) => s.id === columnId);
+
+    if (currentSort && currentSort.desc === desc) {
+      // Clicking the same sort again - remove it
+      clearSorting();
+    } else {
+      // New sort or different direction - move to end of stack
+      const otherSorts = sortingState.filter((s) => s.id !== columnId);
+      const newSort = { id: columnId, desc };
+      setSorting([...otherSorts, newSort]);
+    }
+  };
+
+  it("implements stack-based sorting: moves re-clicked column to end", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: false },
+    ];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age - should move age to end with desc=true
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+    ]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("removes sort when clicking same direction twice", () => {
+    const sortingState: SortingState = [{ id: "age", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Asc on age again - should remove the sort
+    handleSort("age", false, sortingState, setSorting, clearSorting);
+
+    expect(clearSorting).toHaveBeenCalled();
+    expect(setSorting).not.toHaveBeenCalled();
+  });
+
+  it("adds new column to end of stack", () => {
+    const sortingState: SortingState = [{ id: "name", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Asc on age - should add age to end
+    handleSort("age", false, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([
+      { id: "name", desc: false },
+      { id: "age", desc: false },
+    ]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("toggles sort direction when clicking opposite", () => {
+    const sortingState: SortingState = [{ id: "age", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age - should toggle to descending
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([{ id: "age", desc: true }]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("correctly calculates priority numbers", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+      { id: "dept", desc: false },
+    ];
+
+    // Priority is index + 1
+    const nameSort = sortingState.find((s) => s.id === "name");
+    const namePriority = nameSort ? sortingState.indexOf(nameSort) + 1 : null;
+    expect(namePriority).toBe(1);
+
+    const deptSort = sortingState.find((s) => s.id === "dept");
+    const deptPriority = deptSort ? sortingState.indexOf(deptSort) + 1 : null;
+    expect(deptPriority).toBe(3);
+  });
+
+  it("handles removing column from middle of stack", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+      { id: "dept", desc: false },
+    ];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age again - should remove it
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(clearSorting).toHaveBeenCalled();
+    // After removal, dept should move from priority 3 to priority 2
+  });
+});

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

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 "use no memo";
 
-import type { Column } from "@tanstack/react-table";
+import type { Column, Table } from "@tanstack/react-table";
 import { capitalize } from "lodash-es";
 import { FilterIcon, MinusIcon, TextIcon, XIcon } from "lucide-react";
 import { useMemo, useRef, useState } from "react";
@@ -68,13 +68,15 @@ interface DataTableColumnHeaderProps<TData, TValue>
   column: Column<TData, TValue>;
   header: React.ReactNode;
   calculateTopKRows?: CalculateTopKRows;
+  table?: Table<TData>;
 }
 
 export const DataTableColumnHeader = <TData, TValue>({
   column,
   header,
   className,
   calculateTopKRows,
+  table,
 }: DataTableColumnHeaderProps<TData, TValue>) => {
   const [isFilterValueOpen, setIsFilterValueOpen] = useState(false);
   const { locale } = useLocale();
@@ -117,7 +119,7 @@ export const DataTableColumnHeader = <TData, TValue>({
         </DropdownMenuTrigger>
         <DropdownMenuContent align="start">
           {renderDataType(column)}
-          {renderSorts(column)}
+          {renderSorts(column, table)}
           {renderCopyColumn(column)}
           {renderColumnPinning(column)}
           {renderColumnWrapping(column)}

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

@@ -165,7 +165,7 @@ export function generateColumns<T>({
         return row[key as keyof T];
       },
 
-      header: ({ column }) => {
+      header: ({ column, table }) => {
         const stats = chartSpecModel?.getColumnStats(key);
         const dtype = column.columnDef.meta?.dtype;
         const headerTitle = headerTooltip?.[key];
@@ -208,6 +208,7 @@ export function generateColumns<T>({
             header={headerWithTooltip}
             column={column}
             calculateTopKRows={calculateTopKRows}
+            table={table}
           />
         );
 

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

@@ -215,8 +215,13 @@ const DataTableInternal = <TData,>({
     manualPagination: manualPagination,
     getPaginationRowModel: getPaginationRowModel(),
     // sorting
-    ...(setSorting ? { onSortingChange: setSorting } : {}),
+    ...(setSorting
+      ? {
+          onSortingChange: setSorting,
+        }
+      : {}),
     manualSorting: manualSorting,
+    enableMultiSort: true,
     getSortedRowModel: getSortedRowModel(),
     // filtering
     manualFiltering: true,

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

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { PinLeftIcon, PinRightIcon } from "@radix-ui/react-icons";
-import type { Column } from "@tanstack/react-table";
+import type { Column, SortDirection, Table } from "@tanstack/react-table";
 import {
   AlignJustifyIcon,
   ArrowDownWideNarrowIcon,
@@ -163,27 +163,80 @@ export function renderCopyColumn<TData, TValue>(column: Column<TData, TValue>) {
 const AscIcon = ArrowUpNarrowWideIcon;
 const DescIcon = ArrowDownWideNarrowIcon;
 
-export function renderSorts<TData, TValue>(column: Column<TData, TValue>) {
+export function renderSorts<TData, TValue>(
+  column: Column<TData, TValue>,
+  table?: Table<TData>,
+) {
   if (!column.getCanSort()) {
     return null;
   }
 
+  const sortDirection = column.getIsSorted();
+  const sortingIndex = column.getSortIndex();
+
+  const sortingState = table?.getState().sorting;
+  const hasMultiSort = sortingState?.length && sortingState.length > 1;
+
+  const renderSortIndex = () => {
+    return (
+      <span className="ml-auto text-xs font-medium">{sortingIndex + 1}</span>
+    );
+  };
+
+  const renderClearSort = () => {
+    if (!sortDirection) {
+      return null;
+    }
+
+    if (!hasMultiSort) {
+      // render clear sort for this column
+      return (
+        <DropdownMenuItem onClick={() => column.clearSorting()}>
+          <ChevronsUpDown className="mo-dropdown-icon" />
+          Clear sort
+        </DropdownMenuItem>
+      );
+    }
+
+    // render clear sort for all columns
+    return (
+      <DropdownMenuItem onClick={() => table?.resetSorting()}>
+        <ChevronsUpDown className="mo-dropdown-icon" />
+        Clear all sorts
+      </DropdownMenuItem>
+    );
+  };
+
+  const toggleSort = (direction: SortDirection) => {
+    // Clear sort if clicking the same direction
+    if (sortDirection === direction) {
+      column.clearSorting();
+    } else {
+      // Toggle sort direction
+      const descending = direction === "desc";
+      column.toggleSorting(descending, true);
+    }
+  };
+
   return (
     <>
-      <DropdownMenuItem onClick={() => column.toggleSorting(false)}>
+      <DropdownMenuItem
+        onClick={() => toggleSort("asc")}
+        className={sortDirection === "asc" ? "bg-accent" : ""}
+      >
         <AscIcon className="mo-dropdown-icon" />
         Asc
+        {sortDirection === "asc" && renderSortIndex()}
       </DropdownMenuItem>
-      <DropdownMenuItem onClick={() => column.toggleSorting(true)}>
+      <DropdownMenuItem
+        onClick={() => toggleSort("desc")}
+        className={sortDirection === "desc" ? "bg-accent" : ""}
+      >
         <DescIcon className="mo-dropdown-icon" />
         Desc
+        {sortDirection === "desc" && renderSortIndex()}
       </DropdownMenuItem>
-      {column.getIsSorted() && (
-        <DropdownMenuItem onClick={() => column.clearSorting()}>
-          <ChevronsUpDown className="mo-dropdown-icon" />
-          Clear sort
-        </DropdownMenuItem>
-      )}
+      {renderClearSort()}
       <DropdownMenuSeparator />
     </>
   );

frontend/src/plugins/impl/DataTablePlugin.tsx

@@ -205,7 +205,7 @@ type DataTableFunctions = {
     sort?: {
       by: string;
       descending: boolean;
-    };
+    }[];
     query?: string;
     filters?: ConditionType[];
     page_number: number;
@@ -298,7 +298,12 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
       .input(
         z.object({
           sort: z
-            .object({ by: z.string(), descending: z.boolean() })
+            .array(
+              z.object({
+                by: z.string(),
+                descending: z.boolean(),
+              }),
+            )
             .optional(),
           query: z.string().optional(),
           filters: z.array(ConditionSchema).optional(),
@@ -501,19 +506,15 @@ export const LoadingDataTableComponent = memo(
         !props.lazy &&
         !pageSizeChanged;
 
-      if (sorting.length > 1) {
-        Logger.warn("Multiple sort columns are not supported");
-      }
+      // Convert sorting state to API format
+      const sortArgs =
+        sorting.length > 0
+          ? sorting.map((s) => ({ by: s.id, descending: s.desc }))
+          : undefined;
 
       // If we have sort/search/filter, use the search function
       const searchResultsPromise = search<T>({
-        sort:
-          sorting.length > 0
-            ? {
-                by: sorting[0].id,
-                descending: sorting[0].desc,
-              }
-            : undefined,
+        sort: sortArgs,
         query: searchQuery,
         page_number: paginationState.pageIndex,
         page_size: paginationState.pageSize,
@@ -563,16 +564,15 @@ export const LoadingDataTableComponent = memo(
 
     const getRow = useCallback(
       async (rowId: number) => {
+        const sortArgs =
+          sorting.length > 0
+            ? sorting.map((s) => ({ by: s.id, descending: s.desc }))
+            : undefined;
+
         const result = await search<T>({
           page_number: rowId,
           page_size: 1,
-          sort:
-            sorting.length > 0
-              ? {
-                  by: sorting[0].id,
-                  descending: sorting[0].desc,
-                }
-              : undefined,
+          sort: sortArgs,
           query: searchQuery,
           filters: filters.flatMap((filter) => {
             return filterToFilterCondition(

frontend/src/plugins/impl/data-frames/DataFramePlugin.tsx

@@ -63,7 +63,7 @@ type PluginFunctions = {
     sort?: {
       by: string;
       descending: boolean;
-    };
+    }[];
     query?: string;
     filters?: ConditionType[];
     page_number: number;
@@ -117,7 +117,12 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
       .input(
         z.object({
           sort: z
-            .object({ by: z.string(), descending: z.boolean() })
+            .array(
+              z.object({
+                by: z.string(),
+                descending: z.boolean(),
+              }),
+            )
             .optional(),
           query: z.string().optional(),
           filters: z.array(ConditionSchema).optional(),