fix(data-editor): don't crash appending a row to an empty row-oriented editor

marimo/_plugins/ui/_impl/data_editor.py

@@ -296,9 +296,18 @@ def _apply_edits_row_oriented(
     edits: DataEdits,
     schema: nw.Schema | None = None,
 ) -> RowOrientedData:
+    # Capture the columns before applying edits: a `remove` edit can empty `data`
+    # before a later positional edit appends a new row, and that new row must
+    # still carry every original column instead of only the edited one.
+    if data:
+        columns = list(data[0].keys())
+    elif schema is not None:
+        columns = list(schema.names())
+    else:
+        columns = []
     for edit in edits["edits"]:
         if is_positional_edit(edit):
-            _apply_positional_edit_row_oriented(data, edit, schema)
+            _apply_positional_edit_row_oriented(data, edit, schema, columns)
         elif is_row_edit(edit):
             _apply_row_edit_row_oriented(data, edit)
         elif is_column_edit(edit):
@@ -454,13 +463,28 @@ def _apply_positional_edit_row_oriented(
     data: RowOrientedData,
     edit: PositionalEdit,
     schema: nw.Schema | None = None,
+    columns: list[str] | None = None,
 ) -> None:
     """Apply a positional edit to row-oriented data."""
     if edit["rowIdx"] >= len(data):
-        # Create a new row with None values for all columns
-        new_row = {col: None for col in data[0]}
-        data.append(new_row)
-    original_value = data[0][edit["columnId"]] if data else None
+        # Determine the columns for any new row(s): prefer an existing row, then
+        # the columns captured before edits ran (survives a remove-all), then the
+        # schema; always include the edited column.
+        if data:
+            new_columns = list(data[0].keys())
+        elif columns:
+            new_columns = list(columns)
+        elif schema is not None:
+            new_columns = list(schema.names())
+        else:
+            new_columns = []
+        if edit["columnId"] not in new_columns:
+            new_columns.append(edit["columnId"])
+        # Extend through the requested index (mirrors the column-oriented path)
+        # so a non-contiguous rowIdx does not raise IndexError below.
+        while len(data) <= edit["rowIdx"]:
+            data.append({col: None for col in new_columns})
+    original_value = data[0].get(edit["columnId"]) if data else None
     dtype = schema.get(edit["columnId"]) if schema else None
     data[edit["rowIdx"]][edit["columnId"]] = _convert_value(
         edit["value"], original_value, dtype

tests/_plugins/ui/_impl/test_data_editor.py

@@ -118,6 +118,86 @@ def test_apply_edits_new_row():
     ]
 
 
+def test_apply_edits_row_oriented_append_to_empty():
+    # Appending a row to an empty row-oriented editor should not crash, even
+    # though there is no existing row to read column names from.
+    data: list[dict[str, Any]] = []
+    edits = {"edits": [{"rowIdx": 0, "columnId": "A", "value": "x"}]}
+    result = apply_edits(data, edits)
+    assert result == [{"A": "x"}]
+
+
+def test_apply_edits_row_oriented_append_to_empty_with_schema():
+    # When a schema is available, an appended first row should include every
+    # known column, mirroring the non-empty append behavior.
+    data: list[dict[str, Any]] = []
+    edits = {"edits": [{"rowIdx": 0, "columnId": "A", "value": "1"}]}
+    schema = nw.Schema({"A": nw.Int64(), "B": nw.String()})
+    result = apply_edits(data, edits, schema=schema)
+    assert result == [{"A": 1, "B": None}]
+
+
+def test_apply_edits_row_oriented_remove_all_then_add():
+    # Removing every row and then adding one back (an edit replay reachable from
+    # the UI) must append cleanly and keep every original column, not just the
+    # edited one.
+    data = [{"A": 1, "B": "a"}]
+    edits = {
+        "edits": [
+            {"rowIdx": 0, "type": "remove"},
+            {"rowIdx": 0, "columnId": "A", "value": "x"},
+        ]
+    }
+    result = apply_edits(data, edits)
+    assert result == [{"A": "x", "B": None}]
+
+
+def test_apply_edits_row_oriented_remove_all_then_add_multi_column():
+    # The UI emits one positional edit per column; after removing every row,
+    # rebuilding a full row must not drop columns or KeyError on them.
+    data = [{"A": 1, "B": "a"}]
+    edits = {
+        "edits": [
+            {"rowIdx": 0, "type": "remove"},
+            {"rowIdx": 0, "columnId": "A", "value": "x"},
+            {"rowIdx": 0, "columnId": "B", "value": "y"},
+        ]
+    }
+    assert apply_edits(data, edits) == [{"A": "x", "B": "y"}]
+
+
+def test_apply_edits_row_oriented_non_contiguous_row():
+    # A positional edit at a row index past the end must extend the data with
+    # filled rows instead of raising IndexError.
+    data = [{"A": 1, "B": "a"}]
+    edits = {
+        "edits": [
+            {"rowIdx": 0, "type": "remove"},
+            {"rowIdx": 2, "columnId": "A", "value": "x"},
+        ]
+    }
+    assert apply_edits(data, edits) == [
+        {"A": None, "B": None},
+        {"A": None, "B": None},
+        {"A": "x", "B": None},
+    ]
+
+
+def test_data_editor_convert_value_remove_all_then_add_preserves_columns():
+    # End-to-end through the editor's replay path (_convert_value), not just the
+    # helper: removing all rows and adding one back keeps every column.
+    editor = data_editor(data=[{"A": 1, "B": "a"}])
+    result = editor._convert_value(
+        {
+            "edits": [
+                {"rowIdx": 0, "type": "remove"},
+                {"rowIdx": 0, "columnId": "A", "value": "x"},
+            ]
+        }
+    )
+    assert result == [{"A": "x", "B": None}]
+
+
 @pytest.mark.skipif(
     not DependencyManager.polars.has(), reason="Polars not installed"
 )