Add multi-column sorting with stack-based behavior

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

@@ -202,8 +202,33 @@ const DataTableInternal = <TData,>({
     manualPagination: manualPagination,
     getPaginationRowModel: getPaginationRowModel(),
     // sorting
-    ...(setSorting ? { onSortingChange: setSorting } : {}),
+    ...(setSorting
+      ? {
+          onSortingChange: (updaterOrValue) => {
+            // Custom sorting logic - stack behavior
+            const newSorting =
+              typeof updaterOrValue === "function"
+                ? updaterOrValue(sorting || [])
+                : updaterOrValue;
+
+            // Implement stack behavior: if column already exists, remove it and add to front
+            const customSorting = newSorting.reduce(
+              (acc: typeof newSorting, sort) => {
+                // Remove any existing sorts for this column
+                const filtered = acc.filter((s) => s.id !== sort.id);
+                // Add the new sort to the front (most recent)
+                return [sort, ...filtered];
+              },
+              [],
+            );
+
+            setSorting(customSorting);
+          },
+        }
+      : {}),
     manualSorting: manualSorting,
+    enableMultiSort: true,
+    isMultiSortEvent: () => true, // Always enable multi-sort (no shift key required)
     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,
@@ -14,6 +14,7 @@
   ListFilterPlusIcon,
   PinOffIcon,
   WrapTextIcon,
+  XIcon,
 } from "lucide-react";
 import {
   DropdownMenuItem,
@@ -159,25 +160,92 @@
 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;
   }
 
+  // Try to get table from column (TanStack Table should provide this)
+  const tableFromColumn = (column as any).table || table;
+
+  // If table is available (either passed or from column), use full multi-column sort functionality
+  if (tableFromColumn) {
+    const sortingState: SortingState = tableFromColumn.getState().sorting;
+    const currentSort = sortingState.find((s) => s.id === column.id);
+    const sortIndex = currentSort
+      ? sortingState.indexOf(currentSort) + 1
+      : null;
+
+    return (
+      <>
+        <DropdownMenuItem onClick={() => column.toggleSorting(false, true)}>
+          <AscIcon className="mo-dropdown-icon" />
+          Sort Ascending
+          {sortIndex && currentSort && !currentSort.desc && (
+            <span className="ml-auto text-xs bg-blue-100 text-blue-800 px-1 rounded">
+              {sortIndex}
+            </span>
+          )}
+        </DropdownMenuItem>
+        <DropdownMenuItem onClick={() => column.toggleSorting(true, true)}>
+          <DescIcon className="mo-dropdown-icon" />
+          Sort Descending
+          {sortIndex && currentSort && currentSort.desc && (
+            <span className="ml-auto text-xs bg-blue-100 text-blue-800 px-1 rounded">
+              {sortIndex}
+            </span>
+          )}
+        </DropdownMenuItem>
+        {currentSort && (
+          <DropdownMenuItem onClick={() => column.clearSorting()}>
+            <XIcon className="mo-dropdown-icon" />
+            Remove Sort
+            <span className="ml-auto text-xs bg-red-100 text-red-800 px-1 rounded">
+              {sortIndex}
+            </span>
+          </DropdownMenuItem>
+        )}
+        {sortingState.length > 0 && (
+          <DropdownMenuItem onClick={() => tableFromColumn.resetSorting()}>
+            <FilterX className="mo-dropdown-icon" />
+            Clear All Sorts
+          </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)}>
         <AscIcon className="mo-dropdown-icon" />
-        Asc
+        Sort Ascending
+        {isSorted === "asc" && (
+          <span className="ml-auto text-xs bg-blue-100 text-blue-800 px-1 rounded">
+            ✓
+          </span>
+        )}
       </DropdownMenuItem>
-      <DropdownMenuItem onClick={() => column.toggleSorting(true)}>
+      <DropdownMenuItem onClick={() => column.toggleSorting(true, true)}>
         <DescIcon className="mo-dropdown-icon" />
-        Desc
+        Sort Descending
+        {isSorted === "desc" && (
+          <span className="ml-auto text-xs bg-blue-100 text-blue-800 px-1 rounded">
+            ✓
+          </span>
+        )}
       </DropdownMenuItem>
-      {column.getIsSorted() && (
+      {isSorted && (
         <DropdownMenuItem onClick={() => column.clearSorting()}>
-          <ChevronsUpDown className="mo-dropdown-icon" />
-          Clear sort
+          <XIcon className="mo-dropdown-icon" />
+          Remove Sort
         </DropdownMenuItem>
       )}
       <DropdownMenuSeparator />

frontend/src/plugins/impl/DataTablePlugin.tsx

@@ -196,8 +196,8 @@ type DataTableFunctions = {
   ) => Promise<ColumnSummaries<T>>;
   search: <T>(req: {
     sort?: {
-      by: string;
-      descending: boolean;
+      by: string[];
+      descending: boolean[];
     };
     query?: string;
     filters?: ConditionType[];
@@ -282,7 +282,10 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
       .input(
         z.object({
           sort: z
-            .object({ by: z.string(), descending: z.boolean() })
+            .object({
+              by: z.array(z.string()),
+              descending: z.array(z.boolean()),
+            })
             .optional(),
           query: z.string().optional(),
           filters: z.array(ConditionSchema).optional(),
@@ -476,17 +479,13 @@ export const LoadingDataTableComponent = memo(
         !props.lazy &&
         !pageSizeChanged;
 
-      if (sorting.length > 1) {
-        Logger.warn("Multiple sort columns are not supported");
-      }
-
       // 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,
+                by: sorting.map((column) => column.id),
+                descending: sorting.map((column) => column.desc),
               }
             : undefined,
         query: searchQuery,
@@ -541,8 +540,8 @@ export const LoadingDataTableComponent = memo(
           sort:
             sorting.length > 0
               ? {
-                  by: sorting[0].id,
-                  descending: sorting[0].desc,
+                  by: sorting.map((column) => column.id),
+                  descending: sorting.map((column) => column.desc),
                 }
               : undefined,
           query: searchQuery,

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

@@ -56,8 +56,8 @@ type PluginFunctions = {
   }>;
   search: <T>(req: {
     sort?: {
-      by: string;
-      descending: boolean;
+      by: string[];
+      descending: boolean[];
     };
     query?: string;
     filters?: ConditionType[];
@@ -110,7 +110,10 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
       .input(
         z.object({
           sort: z
-            .object({ by: z.string(), descending: z.boolean() })
+            .object({
+              by: z.array(z.string()),
+              descending: z.array(z.boolean()),
+            })
             .optional(),
           query: z.string().optional(),
           filters: z.array(ConditionSchema).optional(),

marimo/_plugins/ui/_impl/dataframes/dataframe.py

@@ -276,8 +276,14 @@ def _apply_filters_query_sort(
         if query:
             result = result.search(query)
 
-        if sort and sort.by in result.get_column_names():
-            result = result.sort_values(sort.by, sort.descending)
+        if sort:
+            # Convert tuples to lists to match the sort_values method signature
+            by_list = list(sort.by)
+            descending_list = list(sort.descending)
+            # Check that all columns exist
+            existing_columns = set(result.get_column_names())
+            if all(col in existing_columns for col in by_list):
+                result = result.sort_values(by_list, descending_list)
 
         return result
 

marimo/_plugins/ui/_impl/table.py

@@ -112,6 +112,12 @@ class ColumnSummaries:
 MaxColumnsType = Union[int, None, MaxColumnsNotProvided]
 
 
+@dataclass(frozen=True)
+class SortArgs:
+    by: tuple[ColumnName, ...]
+    descending: tuple[bool, ...]
+
+
 @dataclass(frozen=True)
 class SearchTableArgs:
     page_size: int
@@ -135,12 +141,6 @@ class SearchTableResponse:
     cell_styles: Optional[CellStyles] = None
 
 
-@dataclass(frozen=True)
-class SortArgs:
-    by: ColumnName
-    descending: bool
-
-
 @dataclass
 class GetRowIdsResponse:
     row_ids: list[int]
@@ -1100,8 +1100,14 @@ def _apply_filters_query_sort(
         if query:
             result = result.search(query)
 
-        if sort and sort.by in result.get_column_names():
-            result = result.sort_values(sort.by, sort.descending)
+        if sort:
+            # Convert tuples to lists to match the sort_values method signature
+            by_list = list(sort.by)
+            descending_list = list(sort.descending)
+            # Check that all columns exist
+            existing_columns = set(result.get_column_names())
+            if all(col in existing_columns for col in by_list):
+                result = result.sort_values(by_list, descending_list)
 
         return result