From 3f36777accd9ae566d7fd63cbace88a6c1f78d73 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Tue, 25 Aug 2026 14:21:44 +0530 Subject: [PATCH 1/2] fix(data-editor): don't crash appending a row to an empty row-oriented editor Appending a positional edit to an empty row-oriented data editor raised IndexError because the new row's columns were read from data[0], which does not exist when there are no rows yet. Derive the columns from the schema when available and always include the edited column, mirroring the column-oriented path. This also fixes the remove-all-then-add edit replay. --- marimo/_plugins/ui/_impl/data_editor.py | 13 ++++++-- tests/_plugins/ui/_impl/test_data_editor.py | 33 +++++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/marimo/_plugins/ui/_impl/data_editor.py b/marimo/_plugins/ui/_impl/data_editor.py index f24574c0408..4dea8827137 100644 --- a/marimo/_plugins/ui/_impl/data_editor.py +++ b/marimo/_plugins/ui/_impl/data_editor.py @@ -457,8 +457,17 @@ def _apply_positional_edit_row_oriented( ) -> 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]} + # Create a new row with None values for all columns. When the editor + # has no rows yet there is no existing row to read column names from, + # so fall back to the schema and always include the edited column. + if data: + columns = list(data[0].keys()) + elif schema is not None: + columns = list(schema.keys()) + else: + columns = [] + new_row: dict[str, Any] = {col: None for col in columns} + new_row.setdefault(edit["columnId"], None) data.append(new_row) original_value = data[0][edit["columnId"]] if data else None dtype = schema.get(edit["columnId"]) if schema else None diff --git a/tests/_plugins/ui/_impl/test_data_editor.py b/tests/_plugins/ui/_impl/test_data_editor.py index 1667183e7d0..9804749dc23 100644 --- a/tests/_plugins/ui/_impl/test_data_editor.py +++ b/tests/_plugins/ui/_impl/test_data_editor.py @@ -118,6 +118,39 @@ 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) should append cleanly instead of raising IndexError. + 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"}] + + @pytest.mark.skipif( not DependencyManager.polars.has(), reason="Polars not installed" ) From 4cd1e136af044084241b1b967d08499a1f0fcde7 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Fri, 28 Aug 2026 09:33:19 +0530 Subject: [PATCH 2/2] Preserve columns and extend rows when replaying positional edits The row-oriented replay only fixed the narrowest empty-append crash and still lost data: after a remove-all, a rebuilt row kept only the edited column; a multi-column rebuild hit KeyError on the second column (it read data[0][col]); and a non-contiguous rowIdx past the end raised IndexError. Capture the columns before edits run (so a remove-all doesn't lose them), thread them into the positional handler, extend the data through the requested index, and read the original value with .get. Also use nw.Schema.names() (the idiomatic narwhals accessor). Fixes the dropped-column / KeyError / IndexError cases and covers them via the helper and end-to-end through _convert_value. --- marimo/_plugins/ui/_impl/data_editor.py | 37 +++++++++----- tests/_plugins/ui/_impl/test_data_editor.py | 53 +++++++++++++++++++-- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/marimo/_plugins/ui/_impl/data_editor.py b/marimo/_plugins/ui/_impl/data_editor.py index 4dea8827137..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,22 +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. When the editor - # has no rows yet there is no existing row to read column names from, - # so fall back to the schema and always include the edited column. + # 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: - columns = list(data[0].keys()) + new_columns = list(data[0].keys()) + elif columns: + new_columns = list(columns) elif schema is not None: - columns = list(schema.keys()) + new_columns = list(schema.names()) else: - columns = [] - new_row: dict[str, Any] = {col: None for col in columns} - new_row.setdefault(edit["columnId"], None) - data.append(new_row) - original_value = data[0][edit["columnId"]] if data else None + 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 9804749dc23..2eabc300518 100644 --- a/tests/_plugins/ui/_impl/test_data_editor.py +++ b/tests/_plugins/ui/_impl/test_data_editor.py @@ -138,8 +138,9 @@ def test_apply_edits_row_oriented_append_to_empty_with_schema(): 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) should append cleanly instead of raising IndexError. + # 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": [ @@ -148,7 +149,53 @@ def test_apply_edits_row_oriented_remove_all_then_add(): ] } result = apply_edits(data, edits) - assert result == [{"A": "x"}] + 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(