From 88203d25156a0dd7fc7d2c51aee66410d546ae18 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:42:36 +0200 Subject: [PATCH] fix(python): keep a renamed Permutation column in its position (#4626) The `Permutation.rename_column` docstring says it "only changes the name of the column that is visible from this permutation". But `permutation.py:868-870` adds the new name at the end of the selection and then deletes the old key, so the renamed column also moves to the last position. Hugging Face `datasets`, which `Permutation` mirrors (#2725), keeps it in place. Positional formats feel it most: `with_format("numpy")` swaps the feature order of every row. Before, on `d97df0c`: ``` $ python repro.py /python/python # lancedb 0.41.0b12 extension, Python sources from the tree ['id', 'value'] ['value', 'row_id'] [{'value': 0, 'row_id': 0}] [[10 1]] ``` After: ``` $ python repro.py /python/python ['id', 'value'] ['row_id', 'value'] [{'row_id': 0, 'value': 0}] [[ 1 10]] ```
repro.py (its first output line, lancedb.__file__, is cut above) ```python import sys; sys.path.insert(0, sys.argv[1]) import tempfile import lancedb from lancedb.permutation import Permutation print(lancedb.__file__) db = lancedb.connect(tempfile.mkdtemp()) tbl = db.create_table("t", [{"id": i, "value": i * 10} for i in range(3)]) perm = Permutation.identity(tbl) print(perm.column_names) renamed = perm.rename_column("id", "row_id") print(renamed.column_names) print(renamed.__getitems__([0])) print(renamed.with_format("numpy").__getitems__([1])) ```
For comparison, `datasets` 5.1.0 gives `['new_id', 'value']` for `Dataset.from_dict({'id': [1, 2], 'value': [3, 4]}).rename_column('id', 'new_id').column_names`. No issue exists for this. The fix rebuilds the selection dict in order. The cost: code that renamed a column and then read it by position in numpy batches (say `[:, -1]`) gets a different column after this, with no error. `rename_columns` calls `rename_column` once per name, so it gets the fix too. The tests that check schema order only rename `value`, the last column, and `test_rename_rowid` checks membership but not order. The new test renames `id`. ### Verification ``` # Python 3.12.6; lancedb 0.41.0b12 wheel from pypi.fury.io, its _lancedb.abi3.so copied into the tree; pylance 14.0.0-rc.1 $ PYTHONPATH=python python -m pytest python/tests/test_permutation.py -q -p no:cacheprovider d97df0c: 72 passed in 2.53s this branch: 73 passed in 2.36s permutation.py reverted, test kept: FAILED python/tests/test_permutation.py::test_rename_column_keeps_position assert value: int64\nnew_id: int64 == new_id: int64\nvalue: int64 1 failed, 72 passed in 2.41s $ ruff format --check . && ruff check . # ruff 0.9.9, as python.yml pins it 152 files already formatted All checks passed! # not run: pyright (its include list covers neither permutation.py nor tests/), the Rust build ``` @westonpace, this is your class from #2725. Do you want the docstring to say outright that the column keeps its position? If so, I'll add that line to both rename methods in this PR. Signed-off-by: Christopher Krah <48338417+chrikrah@users.noreply.github.com> --- python/python/lancedb/permutation.py | 8 +++++--- python/python/tests/test_permutation.py | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 355ebf798..e95db27f6 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -909,9 +909,11 @@ class Permutation: f"Cannot rename column `{old_name}` to `{new_name}` because a column " "with that name already exists" ) - new_selection = self.selection.copy() - new_selection[new_name] = new_selection[old_name] - del new_selection[old_name] + # Rebuild the mapping so the renamed column keeps its position + new_selection = { + (new_name if name == old_name else name): value + for name, value in self.selection.items() + } return self._with_selection(new_selection) def rename_columns(self, column_map: dict[str, str]) -> "Permutation": diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index ed1894340..50ca53061 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation): some_permutation.rename_column("non_existent", "new_value") +def test_rename_column_keeps_position(some_permutation: Permutation): + renamed = some_permutation.rename_column("id", "new_id") + assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())]) + batch = renamed.with_format("arrow").__getitems__([0]) + assert batch.schema.names == ["new_id", "value"] + + def test_rename_columns(some_permutation: Permutation): assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema( [("id", pa.int64()), ("new_value", pa.int64())]