Repository: marimo-team/marimo

Issue:
Select all table rows doesn't select all on filtered table.
### Describe the bug

* Create a `mo.ui.table` from a dataframe
* Filter to a subset of the rows using search
* Select all rows from the top right, then click select all again
* Only a subset of rows are shown to be selected, but all of them should be selected.

See screencast

https://github.com/user-attachments/assets/6a498a71-0d69-45f9-9fab-253c02375980

### Environment

<details>

```
{
  "marimo": "0.11.14",
  "OS": "Darwin",
  "OS Version": "22.5.0",
  "Processor": "arm",
  "Python Version": "3.12.4",
  "Binaries": {
    "Browser": "133.0.6943.142",
    "Node": "v21.5.0"
  },
  "Dependencies": {
    "click": "8.1.8",
    "docutils": "0.21.2",
    "itsdangerous": "2.2.0",
    "jedi": "0.19.2",
    "markdown": "3.7",
    "narwhals": "1.29.0",
    "packaging": "24.2",
    "psutil": "7.0.0",
    "pygments": "2.19.1",
    "pymdown-extensions": "10.14.3",
    "pyyaml": "6.0.2",
    "ruff": "0.9.9",
    "starlette": "0.46.0",
    "tomlkit": "0.13.2",
    "typing-extensions": "4.12.2",
    "uvicorn": "0.34.0",
    "websockets": "15.0"
  },
  "Optional Dependencies": {},
  "Experimental Flags": {
    "reactive_tests": true,
    "rtc": false,
    "multi_column": true,
    "chat_sidebar": true,
    "inline_ai_tooltip": true
  }
}
```

</details>


### Code to reproduce

```python
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "marimo",
#     "vega-datasets==0.9.0",
# ]
# ///

import marimo

__generated_with = "0.11.14"
app = marimo.App(width="medium")


@app.cell
def _():
    from vega_datasets import data
    return (data,)


@app.cell
def _():
    import marimo as mo
    return (mo,)


@app.cell
def _(data, mo):
    df = mo.ui.table(data.cars())
    df
    return (df,)


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()
```

Discussion:
@akshayka: Inspecting `table.value`, it looks like several incorrect rows have been selected (many do not include "chevrolet", despite the filter).

<img width="1119" alt="Image" src="https://github.com/user-attachments/assets/e106ee5b-a989-4d70-978b-473c17ab7220" />@Light2Dark: don't have a straightforward solution atm, but it's likely here: `TableActions.tsx:L59`
```.tsx
 const allKeys = Array.from(
    { length: table.getRowCount() },
    (_, i) => [i, true] as const,
  );
onRowSelectionChange(Object.fromEntries(allKeys));
```

We shouldn't select all like this for filtered tables, because the row_id may not be sequential. 
<details>
<summary>possible solution</summary>
Function on the backend to get the row_ids for filtered tables🤔
</details>@mscolnick: i think for now, we could disallow select-all for correctness when filtering is active 

Solution patch:
diff --git a/frontend/src/components/data-table/TableActions.tsx b/frontend/src/components/data-table/TableActions.tsx
index 235f9ca6..0c16f729 100644
--- a/frontend/src/components/data-table/TableActions.tsx
+++ b/frontend/src/components/data-table/TableActions.tsx
@@ -8,6 +8,8 @@ import { SearchIcon } from "lucide-react";
 import { DataTablePagination } from "./pagination";
 import { DownloadAs, type DownloadActionProps } from "./download-actions";
 import type { Table, RowSelectionState } from "@tanstack/react-table";
+import type { GetRowIds } from "@/plugins/impl/DataTablePlugin";
+import { toast } from "../ui/use-toast";
 
 interface TableActionsProps<TData> {
   enableSearch: boolean;
@@ -20,6 +22,7 @@ interface TableActionsProps<TData> {
   onRowSelectionChange?: (value: RowSelectionState) => void;
   table: Table<TData>;
   downloadAs?: DownloadActionProps["downloadAs"];
+  getRowIds?: GetRowIds;
 }
 
 export const TableActions = <TData,>({
@@ -33,7 +36,52 @@ export const TableActions = <TData,>({
   onRowSelectionChange,
   table,
   downloadAs,
+  getRowIds,
 }: TableActionsProps<TData>) => {
+  const handleSelectAllRows = (value: boolean) => {
+    if (!onRowSelectionChange) {
+      return;
+    }
+
+    // Clear all selections
+    if (!value) {
+      onRowSelectionChange({});
+      return;
+    }
+
+    const selectAllRowsByIndex = () => {
+      const allKeys = Array.from(
+        { length: table.getRowCount() },
+        (_, i) => [i, true] as const,
+      );
+      onRowSelectionChange(Object.fromEntries(allKeys));
+    };
+
+    if (!getRowIds) {
+      selectAllRowsByIndex();
+      return;
+    }
+
+    getRowIds({}).then((data) => {
+      if (data.error) {
+        toast({
+          title: "Not available",
+          description: data.error,
+          variant: "danger",
+        });
+        return;
+      }
+
+      if (data.all_rows) {
+        selectAllRowsByIndex();
+      } else {
+        onRowSelectionChange(
+          Object.fromEntries(data.row_ids.map((id) => [id, true])),
+        );
+      }
+    });
+  };
+
   return (
     <div className="flex items-center justify-between flex-shrink-0 pt-1">
       {onSearchQueryChange && enableSearch && (
@@ -52,21 +100,7 @@ export const TableActions = <TData,>({
         <DataTablePagination
           totalColumns={totalColumns}
           selection={selection}
-          onSelectAllRowsChange={
-            onRowSelectionChange
-              ? (value: boolean) => {
-                  if (value) {
-                    const allKeys = Array.from(
-                      { length: table.getRowCount() },
-                      (_, i) => [i, true] as const,
-                    );
-                    onRowSelectionChange(Object.fromEntries(allKeys));
-                  } else {
-                    onRowSelectionChange({});
-                  }
-                }
-              : undefined
-          }
+          onSelectAllRowsChange={handleSelectAllRows}
           table={table}
         />
       ) : (
diff --git a/frontend/src/components/data-table/data-table.tsx b/frontend/src/components/data-table/data-table.tsx
index 49c49d86..4224c08b 100644
--- a/frontend/src/components/data-table/data-table.tsx
+++ b/frontend/src/components/data-table/data-table.tsx
@@ -30,6 +30,7 @@ import { TableActions } from "./TableActions";
 import { ColumnFormattingFeature } from "./column-formatting/feature";
 import { ColumnWrappingFeature } from "./column-wrapping/feature";
 import { INDEX_COLUMN_NAME } from "./types";
+import type { GetRowIds } from "@/plugins/impl/DataTablePlugin";
 
 interface DataTableProps<TData> extends Partial<DownloadActionProps> {
   wrapperClassName?: string;
@@ -51,6 +52,7 @@ interface DataTableProps<TData> extends Partial<DownloadActionProps> {
   selection?: "single" | "multi" | null;
   rowSelection?: RowSelectionState;
   onRowSelectionChange?: OnChangeFn<RowSelectionState>;
+  getRowIds?: GetRowIds;
   // Search
   enableSearch?: boolean;
   searchQuery?: string;
@@ -82,6 +84,7 @@ const DataTableInternal = <TData,>({
   manualPagination = false,
   pagination = false,
   onRowSelectionChange,
+  getRowIds,
   enableSearch = false,
   searchQuery,
   onSearchQueryChange,
@@ -186,6 +189,7 @@ const DataTableInternal = <TData,>({
         onRowSelectionChange={onRowSelectionChange}
         table={table}
         downloadAs={downloadAs}
+        getRowIds={getRowIds}
       />
     </div>
   );
diff --git a/frontend/src/plugins/impl/DataTablePlugin.tsx b/frontend/src/plugins/impl/DataTablePlugin.tsx
index b4a89c16..89d3d378 100644
--- a/frontend/src/plugins/impl/DataTablePlugin.tsx
+++ b/frontend/src/plugins/impl/DataTablePlugin.tsx
@@ -54,6 +54,12 @@ interface ColumnSummaries<T = unknown> {
   is_disabled?: boolean;
 }
 
+export type GetRowIds = (opts: {}) => Promise<{
+  row_ids: number[];
+  all_rows: boolean;
+  error: string | null;
+}>;
+
 /**
  * Arguments for a data table
  *
@@ -81,7 +87,7 @@ interface Data<T> {
 }
 
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
-type Functions = {
+type DataTableFunctions = {
   download_as: (req: { format: "csv" | "json" }) => Promise<string>;
   get_column_summaries: <T>(opts: {}) => Promise<ColumnSummaries<T>>;
   search: <T>(req: {
@@ -97,6 +103,7 @@ type Functions = {
     data: TableData<T>;
     total_rows: number;
   }>;
+  get_row_ids?: GetRowIds;
 };
 
 type S = Array<string | number>;
@@ -135,7 +142,7 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
       hasStableRowId: z.boolean().default(false),
     }),
   )
-  .withFunctions<Functions>({
+  .withFunctions<DataTableFunctions>({
     download_as: rpc
       .input(z.object({ format: z.enum(["csv", "json"]) }))
       .output(z.string()),
@@ -176,6 +183,13 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
           total_rows: z.number(),
         }),
       ),
+    get_row_ids: rpc.input(z.object({}).passthrough()).output(
+      z.object({
+        row_ids: z.array(z.number()),
+        all_rows: z.boolean(),
+        error: z.string().nullable(),
+      }),
+    ),
   })
   .renderer((props) => {
     return (
@@ -192,7 +206,7 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
     );
   });
 
-interface DataTableProps<T> extends Data<T>, Functions {
+interface DataTableProps<T> extends Data<T>, DataTableFunctions {
   className?: string;
   // Selection
   value: S;
@@ -455,6 +469,7 @@ const DataTableComponent = ({
   textJustifyColumns,
   wrappedColumns,
   totalColumns,
+  get_row_ids,
 }: DataTableProps<unknown> &
   DataTableSearchProps & {
     data: unknown[];
@@ -576,6 +591,7 @@ const DataTableComponent = ({
             onRowSelectionChange={handleRowSelectionChange}
             freezeColumnsLeft={freezeColumnsLeft}
             freezeColumnsRight={freezeColumnsRight}
+            getRowIds={get_row_ids}
           />
         </Labeled>
       </ColumnChartContext.Provider>
diff --git a/marimo/_plugins/ui/_impl/table.py b/marimo/_plugins/ui/_impl/table.py
index 946b87b0..b06a2322 100644
--- a/marimo/_plugins/ui/_impl/table.py
+++ b/marimo/_plugins/ui/_impl/table.py
@@ -103,6 +103,13 @@ class SortArgs:
     descending: bool
 
 
+@dataclass
+class GetRowIdsResponse:
+    row_ids: list[int]
+    all_rows: bool
+    error: Optional[str] = None
+
+
 @mddoc
 class table(
     UIElement[
@@ -416,6 +423,11 @@ class table(
                     arg_cls=SearchTableArgs,
                     function=self._search,
                 ),
+                Function(
+                    name="get_row_ids",
+                    arg_cls=EmptyArgs,
+                    function=self._get_row_ids,
+                ),
             ),
         )
 
@@ -649,6 +661,49 @@ class table(
             total_rows=result.get_num_rows(force=True) or 0,
         )
 
+    def _get_row_ids(self, args: EmptyArgs) -> GetRowIdsResponse:
+        """Get row IDs of a table. If searched, return searched rows else all rows."""
+        del args
+
+        total_rows = self._manager.get_num_rows()
+        num_rows_searched = self._searched_manager.get_num_rows()
+
+        if total_rows is None or num_rows_searched is None:
+            return GetRowIdsResponse(
+                row_ids=[],
+                all_rows=False,
+                error="Failed to get row IDs: number of rows is unknown",
+            )
+
+        # If no search has been applied, return with all_rows=True to avoid passing row IDs
+        if total_rows == num_rows_searched:
+            return GetRowIdsResponse(row_ids=[], all_rows=True)
+
+        if num_rows_searched > 1_000_000:
+            return GetRowIdsResponse(
+                row_ids=[],
+                all_rows=False,
+                error="Select all with search is not supported for large datasets. Please filter to less than 1,000,000 rows",
+            )
+
+        # For dictionary data, return sequential indices
+        if isinstance(self.data, dict):
+            return GetRowIdsResponse(
+                row_ids=list(range(num_rows_searched)),
+                all_rows=False,
+            )
+
+        # For dataframes
+        try:
+            row_ids = self._searched_manager.data[INDEX_COLUMN_NAME].to_list()
+            return GetRowIdsResponse(row_ids=row_ids, all_rows=False)
+        except Exception as e:
+            return GetRowIdsResponse(
+                row_ids=[],
+                all_rows=False,
+                error=f"Failed to get row IDs: {str(e)}",
+            )
+
     def _repr_markdown_(self) -> str:
         """Return a markdown representation of the table.
 


Test patch:
diff --git a/tests/_plugins/ui/_impl/test_table.py b/tests/_plugins/ui/_impl/test_table.py
index 9f8eb128..245864c1 100644
--- a/tests/_plugins/ui/_impl/test_table.py
+++ b/tests/_plugins/ui/_impl/test_table.py
@@ -444,6 +444,82 @@ def test_search_sort_nonexistent_columns() -> None:
     assert table._convert_value(["0"]) == ["banana"]
 
 
+def test_get_row_ids() -> None:
+    data = {
+        "id": [1, 2, 3] * 3,
+        "fruits": ["banana", "apple", "cherry"] * 3,
+        "quantity": [10, 20, 30] * 3,
+    }
+    table = ui.table(data)
+
+    initial_response = table._get_row_ids(EmptyArgs())
+    assert initial_response.all_rows is True
+    assert initial_response.row_ids == []
+    assert initial_response.error is None
+
+    table._search(
+        SearchTableArgs(
+            query="cherry",
+            page_size=10,
+            page_number=0,
+        )
+    )
+
+    response = table._get_row_ids(EmptyArgs())
+    # For dicts, we do not need to find row_id, we just return the index
+    assert response.row_ids == [0, 1, 2]
+    assert response.all_rows is False
+    assert response.error is None
+
+
+@pytest.mark.parametrize(
+    "df",
+    create_dataframes(
+        {
+            "id": [1, 2, 3] * 3,
+            "fruits": ["banana", "apple", "cherry"] * 3,
+            "quantity": [10, 20, 30] * 3,
+        },
+        exclude=["ibis", "duckdb", "pyarrow"],
+    ),
+)
+def test_get_row_ids_with_df(df: any) -> None:
+    table = ui.table(df)
+
+    initial_response = table._get_row_ids(EmptyArgs())
+    assert initial_response.all_rows is True
+    assert initial_response.row_ids == []
+    assert initial_response.error is None
+
+    # Test with search
+    table._search(
+        SearchTableArgs(
+            query="cherry",
+            page_size=10,
+            page_number=0,
+        )
+    )
+
+    response = table._get_row_ids(EmptyArgs())
+    assert response.row_ids == [2, 5, 8]
+    assert response.all_rows is False
+    assert response.error is None
+
+    # Test with no search
+    table._search(
+        SearchTableArgs(
+            query="",
+            page_size=10,
+            page_number=0,
+        )
+    )
+
+    response = table._get_row_ids(EmptyArgs())
+    assert response.all_rows is True
+    assert response.row_ids == []
+    assert response.error is None
+
+
 def test_table_with_too_many_columns_passes() -> None:
     data = {str(i): [1] for i in range(101)}
     assert ui.table(data) is not None