Repository: marimo-team/marimo
Pull request: fix: preserve table selection on accordion toggle
Fix table selection being cleared when placed in accordion.

- Modified useEffect to only clear selection when filters/search/sort change
- Added test for selection state persistence
- Fixes #3142
- Note: pre-commit.ci failures are in unrelated files (openapi/src/api.ts, combobox.tsx, cell-link.tsx) that weren't modified by this PR
- Note: CLI test failures appear to be unrelated to frontend changes

Link to Devin run: https://app.devin.ai/sessions/570ca38070d646fb8b5569464068d600

Problem statement:
table selection is cleared when placed in accordion and this is folded / unfolded
### Describe the bug

I have a table that I want to select a row from. The choice of row is used to populate a text field below, as part of a configuration. 

When I place the table inside an accordion, or another element, in which it is refreshed upon folding / unfolding, the selection is scrapped and my text field is empty again.

### Environment

{
  "marimo": "0.9.34",
  "OS": "Windows",
  "OS Version": "11",
  "Processor": "Intel64 Family 6 Model 126 Stepping 5, GenuineIntel",
  "Python Version": "3.12.2",
  "Binaries": {
    "Browser": "129.0.6668.60",
    "Node": "v20.14.0"
  },
  "Dependencies": {
    "click": "8.1.7",
    "docutils": "0.20.1",
    "itsdangerous": "2.2.0",
    "jedi": "0.19.1",
    "markdown": "3.6",
    "narwhals": "1.14.0",
    "packaging": "24.0",
    "psutil": "5.9.8",
    "pygments": "2.17.2",
    "pymdown-extensions": "10.7.1",
    "pyyaml": "6.0.1",
    "ruff": "0.5.2",
    "starlette": "0.37.2",
    "tomlkit": "0.12.4",
    "typing-extensions": "4.12.2",
    "uvicorn": "0.29.0",
    "websockets": "11.0.3"
  },
  "Optional Dependencies": {
    "pandas": "2.2.2"
  }
}


### Code to reproduce

Here's a way to test it. Unfold the accordion, select a row,  see how the text field is updated, fold the accordion, and unfold it, and see how the selection disappears, and the text field is now empty.

```python

import marimo

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


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


@app.cell
def __(mo):
    table=mo.ui.table([{"col1": "hello"},{"col1": "world"}], selection="single")
    return (table,)


@app.cell
def __(mo, table):
    if table.value:
        val = table.value[0]["col1"]
    else:
        val = ""
    text=mo.ui.text(label="Select from table", value=val)

    return text, val


@app.cell
def __(mo, table, text):
    mo.accordion({"fold down": mo.vstack([table, text])})
    return


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

Hints from issue discussion:
Seems like a react state/initialization bug. Thanks for pointing this out.

Commit message: refactor: rename useSkipFirstRender to useEffectSkipFirstRender

Co-Authored-By: Myles Scolnick <myles@marimo.io>

Patch to review:
diff --git a/frontend/src/components/data-table/__tests__/data-table.test.tsx b/frontend/src/components/data-table/__tests__/data-table.test.tsx
new file mode 100644
index 00000000000..6e78d53db04
--- /dev/null
+++ b/frontend/src/components/data-table/__tests__/data-table.test.tsx
@@ -0,0 +1,52 @@
+import { render } from "@testing-library/react";
+import { DataTable } from "../data-table";
+import { describe, expect, it } from "vitest";
+import { vi } from "vitest";
+import type { ColumnDef } from "@tanstack/react-table";
+import type { RowSelectionState } from "@tanstack/react-table";
+
+interface TestData {
+  id: number;
+  name: string;
+}
+
+describe("DataTable", () => {
+  it("should maintain selection state when remounted", () => {
+    const mockOnRowSelectionChange = vi.fn();
+    const testData: TestData[] = [
+      { id: 1, name: "Test 1" },
+      { id: 2, name: "Test 2" },
+    ];
+
+    const columns: Array<ColumnDef<TestData>> = [
+      { accessorKey: "name", header: "Name" },
+    ];
+
+    const initialRowSelection: RowSelectionState = { "0": true };
+
+    const commonProps = {
+      data: testData,
+      columns,
+      selection: "single" as const,
+      totalRows: 2,
+      totalColumns: 1,
+      pagination: false,
+      rowSelection: initialRowSelection,
+      onRowSelectionChange: mockOnRowSelectionChange,
+    };
+
+    const { rerender } = render(<DataTable {...commonProps} />);
+
+    // Verify initial selection is not cleared
+    expect(mockOnRowSelectionChange).not.toHaveBeenCalledWith({});
+
+    // Simulate remount (as would happen in accordion toggle)
+    rerender(<DataTable {...commonProps} />);
+
+    // Verify selection is still not cleared after remount
+    expect(mockOnRowSelectionChange).not.toHaveBeenCalledWith({});
+
+    // Verify the rowSelection prop is maintained
+    expect(commonProps.rowSelection).toEqual(initialRowSelection);
+  });
+});
diff --git a/frontend/src/hooks/__tests__/useEffectSkipFirstRender.test.tsx b/frontend/src/hooks/__tests__/useEffectSkipFirstRender.test.tsx
new file mode 100644
index 00000000000..e45e07b18c7
--- /dev/null
+++ b/frontend/src/hooks/__tests__/useEffectSkipFirstRender.test.tsx
@@ -0,0 +1,41 @@
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { useEffectSkipFirstRender } from "../useEffectSkipFirstRender";
+
+describe("useEffectSkipFirstRender", () => {
+  it("should skip callback on first render", () => {
+    const callback = vi.fn();
+    const { rerender } = renderHook(
+      ({ deps }) => useEffectSkipFirstRender(callback, deps),
+      { initialProps: { deps: [1] } }
+    );
+
+    expect(callback).not.toHaveBeenCalled();
+
+    rerender({ deps: [2] });
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
+  it("should call callback on subsequent dependency changes", () => {
+    const callback = vi.fn();
+    const { rerender } = renderHook(
+      ({ deps }) => useEffectSkipFirstRender(callback, deps),
+      { initialProps: { deps: [1] } }
+    );
+
+    rerender({ deps: [2] });
+    rerender({ deps: [3] });
+    expect(callback).toHaveBeenCalledTimes(2);
+  });
+
+  it("should not call callback if dependencies haven't changed", () => {
+    const callback = vi.fn();
+    const { rerender } = renderHook(
+      ({ deps }) => useEffectSkipFirstRender(callback, deps),
+      { initialProps: { deps: [1] } }
+    );
+
+    rerender({ deps: [1] });
+    expect(callback).not.toHaveBeenCalled();
+  });
+});
diff --git a/frontend/src/hooks/useEffectSkipFirstRender.ts b/frontend/src/hooks/useEffectSkipFirstRender.ts
new file mode 100644
index 00000000000..3ed2a070cca
--- /dev/null
+++ b/frontend/src/hooks/useEffectSkipFirstRender.ts
@@ -0,0 +1,20 @@
+import { useEffect, useRef } from "react";
+
+/**
+ * A hook that skips running a callback on the first render,
+ * then executes it on subsequent renders when dependencies change.
+ *
+ * @param callback - Function to execute after first render
+ * @param deps - Dependencies that trigger the callback when changed
+ */
+export function useEffectSkipFirstRender(callback: () => void, deps: React.DependencyList) {
+  const hasRendered = useRef(false);
+
+  useEffect(() => {
+    if (!hasRendered.current) {
+      hasRendered.current = true;
+      return;
+    }
+    callback();
+  }, deps);
+}
diff --git a/frontend/src/plugins/impl/DataTablePlugin.tsx b/frontend/src/plugins/impl/DataTablePlugin.tsx
index 0fe2115afc9..d0aae7ecea9 100644
--- a/frontend/src/plugins/impl/DataTablePlugin.tsx
+++ b/frontend/src/plugins/impl/DataTablePlugin.tsx
@@ -44,6 +44,7 @@ import { useAsyncData } from "@/hooks/useAsyncData";
 import { useDeepCompareMemoize } from "@/hooks/useDeepCompareMemoize";
 import { DelayMount } from "@/components/utils/delay-mount";
 import { DATA_TYPES } from "@/core/kernel/messages";
+import { useEffectSkipFirstRender } from "../../hooks/useEffectSkipFirstRender";
 
 type CsvURL = string;
 type TableData<T> = T[] | CsvURL;
@@ -232,10 +233,12 @@ export const LoadingDataTableComponent = memo(
 
     // We need to clear the selection when sort, query, or filters change
     // Currently, our selection is index-based,
-    // so we can't rely on the data to be the same
+    // so we can't rely on the data to be the same when these change
     // We can remove this when we have a stable key for each row
-    useEffect(() => {
-      setValue([]);
+    useEffectSkipFirstRender(() => {
+      if (filters.length > 0 || searchQuery !== "" || sorting.length > 0) {
+        setValue([]);
+      }
     }, [setValue, filters, searchQuery, sorting]);
 
     // If pageSize changes, reset pagination state


Code review comments:
diff_hunk: @@ -232,10 +233,12 @@ export const LoadingDataTableComponent = memo(
 
     // We need to clear the selection when sort, query, or filters change
     // Currently, our selection is index-based,
-    // so we can't rely on the data to be the same
+    // so we can't rely on the data to be the same when these change
     // We can remove this when we have a stable key for each row
-    useEffect(() => {
-      setValue([]);
+    useEffectSkipFirstRender(() => {
+      if (filters.length > 0 || searchQuery !== "" || sorting.length > 0) {
+        setValue([]);
+      }
path: frontend/src/plugins/impl/DataTablePlugin.tsx
text: @user1:
```suggestion
      setValue([]);
```