Add multi-column sorting with stack-based behavior

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, SortingState, Table } from "@tanstack/react-table";
 import {
   AlignJustifyIcon,
   ArrowDownWideNarrowIcon,
@@ -163,22 +163,99 @@ 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;
   }
 
+  // If table is available, use full multi-column sort functionality
+  if (table) {
+    const sortingState: SortingState = table.getState().sorting;
+    const currentSort = sortingState.find((s) => s.id === column.id);
+    const sortIndex = currentSort
+      ? sortingState.indexOf(currentSort) + 1
+      : null;
+
+    // Handler to implement stack-based sorting: clicking a sort moves it to the end (highest priority)
+    // Clicking the same sort direction again removes it
+    const handleSort = (desc: boolean) => {
+      if (currentSort && currentSort.desc === desc) {
+        // Clicking the same sort again - remove it
+        column.clearSorting();
+      } else {
+        // New sort or different direction - move to end of stack
+        const otherSorts = sortingState.filter((s) => s.id !== column.id);
+        const newSort = { id: column.id, desc };
+        table.setSorting([...otherSorts, newSort]);
+      }
+    };
+
+    const isActiveSort = (desc: boolean) =>
+      sortIndex && currentSort && currentSort.desc === desc;
+
+    return (
+      <>
+        <DropdownMenuItem
+          onClick={() => handleSort(false)}
+          className={isActiveSort(false) ? "bg-accent" : ""}
+        >
+          <AscIcon className="mo-dropdown-icon" />
+          Asc
+          {isActiveSort(false) && (
+            <span className="ml-auto text-xs font-medium">{sortIndex}</span>
+          )}
+        </DropdownMenuItem>
+        <DropdownMenuItem
+          onClick={() => handleSort(true)}
+          className={isActiveSort(true) ? "bg-accent" : ""}
+        >
+          <DescIcon className="mo-dropdown-icon" />
+          Desc
+          {isActiveSort(true) && (
+            <span className="ml-auto text-xs font-medium">{sortIndex}</span>
+          )}
+        </DropdownMenuItem>
+        {sortingState.length > 1 ? (
+          <DropdownMenuItem onClick={() => table.resetSorting()}>
+            <ChevronsUpDown className="mo-dropdown-icon" />
+            Clear all sorts
+          </DropdownMenuItem>
+        ) : (
+          currentSort && (
+            <DropdownMenuItem onClick={() => column.clearSorting()}>
+              <ChevronsUpDown className="mo-dropdown-icon" />
+              Clear sort
+            </DropdownMenuItem>
+          )
+        )}
+        <DropdownMenuSeparator />
+      </>
+    );
+  }
+
+  // Fallback to simple sorting if table not provided
+  const isSorted = column.getIsSorted();
+
   return (
     <>
-      <DropdownMenuItem onClick={() => column.toggleSorting(false)}>
+      <DropdownMenuItem
+        onClick={() => column.toggleSorting(false, true)}
+        className={isSorted === "asc" ? "bg-accent" : ""}
+      >
         <AscIcon className="mo-dropdown-icon" />
         Asc
       </DropdownMenuItem>
-      <DropdownMenuItem onClick={() => column.toggleSorting(true)}>
+      <DropdownMenuItem
+        onClick={() => column.toggleSorting(true, true)}
+        className={isSorted === "desc" ? "bg-accent" : ""}
+      >
         <DescIcon className="mo-dropdown-icon" />
         Desc
       </DropdownMenuItem>
-      {column.getIsSorted() && (
+      {isSorted && (
         <DropdownMenuItem onClick={() => column.clearSorting()}>
           <ChevronsUpDown className="mo-dropdown-icon" />
           Clear sort

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(),