mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
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 <tree>/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 <tree>/python/python
['id', 'value']
['row_id', 'value']
[{'row_id': 0, 'value': 0}]
[[ 1 10]]
```
<details><summary>repro.py (its first output line,
<code>lancedb.__file__</code>, is cut above)</summary>
```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]))
```
</details>
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>
This commit is contained in:
@@ -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":
|
||||
|
||||
@@ -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())]
|
||||
|
||||
Reference in New Issue
Block a user