fix(data-editor): make edit replay robust across structural changes

frontend/src/plugins/impl/DataEditorPlugin.tsx

@@ -8,6 +8,7 @@ import { LoadingTable } from "@/components/data-table/loading-table";
 import { type FieldTypes, toFieldTypes } from "@/components/data-table/types";
 import { Alert, AlertTitle } from "@/components/ui/alert";
 import { DelayMount } from "@/components/utils/delay-mount";
+import type { DataType } from "@/core/kernel/messages";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { createPlugin } from "../core/builder";
 import type { Setter } from "../types";
@@ -17,6 +18,7 @@ import {
   type DataEditorProps,
   type Edits,
 } from "./data-editor/types";
+import { orderColumnFields } from "./data-editor/data-utils";
 import { vegaLoadData } from "./vega/loader";
 import { getVegaFieldTypes } from "./vega/utils";
 
@@ -45,6 +47,7 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
       label: z.string().nullable(),
       data: z.union([z.string(), z.array(z.object({}).passthrough())]),
       fieldTypes: columnToFieldTypesSchema.nullish(),
+      columnNames: z.array(z.string()).default([]),
       editableColumns: z.union([z.array(z.string()), z.literal("all")]),
       columnSizingMode: z.enum(["auto", "fit"]).default("auto"), // TODO: Remove this
     }),
@@ -55,6 +58,7 @@ export const DataEditorPlugin = createPlugin<Edits>("marimo-data-editor", {
       <LoadingDataEditor
         data={props.data.data}
         fieldTypes={props.data.fieldTypes}
+        columnNames={props.data.columnNames}
         edits={props.value}
         onEdits={props.setValue}
         host={props.host}
@@ -72,10 +76,11 @@ interface Props extends Omit<
   onEdits: Setter<Edits>;
   host: HTMLElement;
   editableColumns: string[] | "all";
+  columnNames: string[];
 }
 
 const LoadingDataEditor = (props: Props) => {
-  const [data, setData] = useState<unknown[]>([]);
+  const [data, setData] = useState<unknown[] | null>(null);
   const [columnFields, setColumnFields] = useState<FieldTypes>(new Map());
 
   // Load the data
@@ -98,9 +103,12 @@ const LoadingDataEditor = (props: Props) => {
 
     setData(localData);
     setColumnFields(
-      toFieldTypes(props.fieldTypes ?? inferFieldTypes(localData)),
+      orderColumnFields(
+        toFieldTypes(props.fieldTypes ?? inferFieldTypes(localData)),
+        props.columnNames,
+      ),
     );
-  }, [props.fieldTypes, props.data]);
+  }, [props.fieldTypes, props.columnNames, props.data]);
 
   if (error) {
     return (
@@ -113,7 +121,7 @@ const LoadingDataEditor = (props: Props) => {
     );
   }
 
-  if (!data) {
+  if (data === null) {
     return (
       <DelayMount milliseconds={200}>
         <LoadingTable pageSize={10} />
@@ -124,7 +132,13 @@ const LoadingDataEditor = (props: Props) => {
   return (
     <LazyDataEditor
       data={data}
-      setData={setData}
+      setData={(nextData) => {
+        if (typeof nextData === "function") {
+          setData((previousData) => nextData(previousData ?? []));
+        } else {
+          setData(nextData);
+        }
+      }}
       columnFields={columnFields}
       setColumnFields={setColumnFields}
       editableColumns={props.editableColumns}
@@ -166,10 +180,13 @@ const LoadingDataEditor = (props: Props) => {
           edits: [...v.edits, { columnIdx, type: BulkEdit.Remove }],
         }));
       }}
-      onAddColumn={(columnIdx: number, newName: string) => {
+      onAddColumn={(columnIdx: number, newName: string, dataType: DataType) => {
         props.onEdits((v) => ({
           ...v,
-          edits: [...v.edits, { columnIdx, newName, type: BulkEdit.Insert }],
+          edits: [
+            ...v.edits,
+            { columnIdx, newName, dataType, type: BulkEdit.Insert },
+          ],
         }));
       }}
     />

frontend/src/plugins/impl/data-editor/__tests__/data-utils.test.ts

@@ -6,6 +6,7 @@ import type { DataType } from "@/core/kernel/messages";
 import {
   insertColumn,
   modifyColumnFields,
+  orderColumnFields,
   removeColumn,
   renameColumn,
 } from "../data-utils";
@@ -15,6 +16,26 @@ import {
 const asFieldTypes = (obj: Record<string, DataType>): FieldTypes =>
   new Map(Object.entries(obj));
 
+describe("orderColumnFields", () => {
+  it("uses explicit order for numeric-looking column names", () => {
+    const inferred = new Map([
+      ["2", "string"],
+      ["10", "string"],
+    ]) as FieldTypes;
+
+    expect([...orderColumnFields(inferred, ["10", "2"])]).toEqual([
+      ["10", "string"],
+      ["2", "string"],
+    ]);
+  });
+
+  it("creates unknown fields for an empty table", () => {
+    expect([...orderColumnFields(new Map(), ["value"])]).toEqual([
+      ["value", "unknown"],
+    ]);
+  });
+});
+
 describe("removeColumn", () => {
   const testData = [
     { int: 1, string: "a", bool: "True", datetime: "2025-07-12 00:07:13" },
@@ -23,8 +44,8 @@ describe("removeColumn", () => {
     { int: 0, string: "", bool: "", datetime: "" },
   ];
 
-  it("should remove column at index 0", () => {
-    const result = removeColumn(testData, 0);
+  it("should remove the first column", () => {
+    const result = removeColumn(testData, "int");
 
     expect(result).toEqual([
       { string: "a", bool: "True", datetime: "2025-07-12 00:07:13" },
@@ -34,8 +55,8 @@ describe("removeColumn", () => {
     ]);
   });
 
-  it("should remove column at index 1", () => {
-    const result = removeColumn(testData, 1);
+  it("should remove the second column", () => {
+    const result = removeColumn(testData, "string");
 
     expect(result).toEqual([
       { int: 1, bool: "True", datetime: "2025-07-12 00:07:13" },
@@ -45,8 +66,8 @@ describe("removeColumn", () => {
     ]);
   });
 
-  it("should remove column at index 2", () => {
-    const result = removeColumn(testData, 2);
+  it("should remove the third column", () => {
+    const result = removeColumn(testData, "bool");
 
     expect(result).toEqual([
       { int: 1, string: "a", datetime: "2025-07-12 00:07:13" },
@@ -56,8 +77,8 @@ describe("removeColumn", () => {
     ]);
   });
 
-  it("should remove column at index 3", () => {
-    const result = removeColumn(testData, 3);
+  it("should remove the last column", () => {
+    const result = removeColumn(testData, "datetime");
 
     expect(result).toEqual([
       { int: 1, string: "a", bool: "True" },
@@ -67,29 +88,22 @@ describe("removeColumn", () => {
     ]);
   });
 
-  it("should handle removing non-existent column index", () => {
-    const result = removeColumn(testData, 999);
-    // Should return the original data since the index doesn't exist
-    expect(result).toEqual(testData);
-  });
-
-  it("should handle negative index", () => {
-    const result = removeColumn(testData, -1);
-    // Should return the original data since negative index is invalid
+  it("should handle removing a non-existent column", () => {
+    const result = removeColumn(testData, "missing");
     expect(result).toEqual(testData);
   });
 
   it("should preserve original data structure", () => {
     const originalData = [...testData];
-    removeColumn(testData, 1);
+    removeColumn(testData, "string");
     // Original data should remain unchanged
     expect(testData).toEqual(originalData);
   });
 
   it("should handle objects with different property counts", () => {
     const mixedData = [{ a: 1, b: 2 }, { a: 3, b: 4, c: 5 }, { a: 6 }];
 
-    const result = removeColumn(mixedData, 1);
+    const result = removeColumn(mixedData, "b");
 
     expect(result).toEqual([{ a: 1 }, { a: 3, c: 5 }, { a: 6 }]);
   });
@@ -100,7 +114,7 @@ describe("removeColumn", () => {
       { a: null, b: 2, c: 3 },
     ];
 
-    const result = removeColumn(dataWithNulls, 1);
+    const result = removeColumn(dataWithNulls, "b");
 
     expect(result).toEqual([
       { a: 1, c: undefined },
@@ -118,7 +132,7 @@ describe("insertColumn", () => {
   ];
 
   it("should insert column at index 0", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 0);
 
     const expected = [
       {
@@ -139,10 +153,17 @@ describe("insertColumn", () => {
       { newColumn: "", int: 0, string: "", bool: "", datetime: "" },
     ];
     expect(result).toEqual(expected);
+    expect(Object.keys(result[0])).toEqual([
+      "newColumn",
+      "int",
+      "string",
+      "bool",
+      "datetime",
+    ]);
   });
 
   it("should insert column at index 1", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 1);
 
     expect(result).toEqual([
       {
@@ -162,10 +183,17 @@ describe("insertColumn", () => {
       },
       { int: 0, newColumn: "", string: "", bool: "", datetime: "" },
     ]);
+    expect(Object.keys(result[0])).toEqual([
+      "int",
+      "newColumn",
+      "string",
+      "bool",
+      "datetime",
+    ]);
   });
 
   it("should insert column at index 2", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 2);
 
     expect(result).toEqual([
       {
@@ -188,7 +216,7 @@ describe("insertColumn", () => {
   });
 
   it("should insert column at index 3", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 3);
 
     expect(result).toEqual([
       {
@@ -211,7 +239,7 @@ describe("insertColumn", () => {
   });
 
   it("should insert column at the end when index equals length", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 4);
 
     expect(result).toEqual([
       {
@@ -234,7 +262,7 @@ describe("insertColumn", () => {
   });
 
   it("should handle inserting at index beyond array length", () => {
-    const result = insertColumn(testData, "newColumn");
+    const result = insertColumn(testData, "newColumn", 999);
     // Should add the column at the end
     expect(result).toEqual([
       {
@@ -270,14 +298,14 @@ describe("insertColumn", () => {
 
   it("should handle array with single object", () => {
     const singleRowData = [{ a: 1, b: 2, c: 3 }];
-    const result = insertColumn(singleRowData, "newColumn");
+    const result = insertColumn(singleRowData, "newColumn", 1);
     expect(result).toEqual([{ a: 1, newColumn: "", b: 2, c: 3 }]);
   });
 
   it("should handle objects with different property counts", () => {
     const mixedData = [{ a: 1, b: 2 }, { a: 3, b: 4, c: 5 }, { a: 6 }];
 
-    const result = insertColumn(mixedData, "newColumn");
+    const result = insertColumn(mixedData, "newColumn", 1);
 
     expect(result).toEqual([
       { a: 1, newColumn: "", b: 2 },
@@ -292,7 +320,7 @@ describe("insertColumn", () => {
       { a: null, b: 2, c: 3 },
     ];
 
-    const result = insertColumn(dataWithNulls, "newColumn");
+    const result = insertColumn(dataWithNulls, "newColumn", 1);
 
     expect(result).toEqual([
       { a: 1, newColumn: "", b: null, c: undefined },
@@ -301,7 +329,7 @@ describe("insertColumn", () => {
   });
 
   it("should handle special characters in column name", () => {
-    const result = insertColumn(testData, "new-column_with_123");
+    const result = insertColumn(testData, "new-column_with_123", 1);
 
     expect(result).toEqual([
       {
@@ -373,12 +401,7 @@ describe("renameColumn", () => {
   it("should handle renaming to an existing column name", () => {
     const result = renameColumn(testData, "int", "string");
 
-    expect(result).toEqual([
-      { string: 1, bool: "True", datetime: "2025-07-12 00:07:13" },
-      { string: 2, bool: "False", datetime: null },
-      { string: 3, bool: null, datetime: "2025-07-12 00:07:13" },
-      { string: 0, bool: "", datetime: "" },
-    ]);
+    expect(result).toEqual(testData);
   });
 
   it("should handle non-existent column name gracefully", () => {

frontend/src/plugins/impl/data-editor/__tests__/glide-data-editor.test.tsx

@@ -57,6 +57,7 @@ describe("GlideDataEditor portal", () => {
   let fullscreenElement: Element | null;
 
   beforeEach(() => {
+    vi.clearAllMocks();
     fullscreenElement = null;
     Object.defineProperty(document, "fullscreenElement", {
       get: () => fullscreenElement,
@@ -85,6 +86,24 @@ describe("GlideDataEditor portal", () => {
     });
   });
 
+  it("replays initial edits when the source data is empty", async () => {
+    const setData = vi.fn();
+    render(
+      <TooltipProvider>
+        <GlideDataEditor
+          {...editorProps}
+          data={[]}
+          setData={setData}
+          edits={[{ rowIdx: 0, columnId: "name", value: "alice" }]}
+        />
+      </TooltipProvider>,
+    );
+
+    await waitFor(() => {
+      expect(setData).toHaveBeenCalledWith([{ name: "alice" }]);
+    });
+  });
+
   it("mounts into the fullscreen element when fullscreen is already active", () => {
     const fullscreenContainer = document.createElement("div");
     document.body.appendChild(fullscreenContainer);