diff --git a/marimo/_plugins/ui/_impl/data_editor.py b/marimo/_plugins/ui/_impl/data_editor.py index f24574c0408..52dc45efbe7 100644 --- a/marimo/_plugins/ui/_impl/data_editor.py +++ b/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 diff --git a/tests/_plugins/ui/_impl/test_data_editor.py b/tests/_plugins/ui/_impl/test_data_editor.py index 1667183e7d0..2eabc300518 100644 --- a/tests/_plugins/ui/_impl/test_data_editor.py +++ b/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" )