diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 6f585ad76..5f93bb7dd 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2183,7 +2183,7 @@ class LanceTable(Table): ``create``/``checkout`` return a new table handle scoped to the branch; writes on it do not affect ``main``. """ - return Branches(self._table) + return Branches(self) def checkout(self, version: Union[int, str]): """Checkout a version of the table. This is an in-place operation. @@ -5809,8 +5809,9 @@ class Branches: Table branch manager. """ - def __init__(self, table): - self._table = table + def __init__(self, parent: "LanceTable"): + self._parent = parent + self._table = parent._table def list(self) -> Dict[str, Any]: """List all branches, mapping name to branch metadata.""" @@ -5836,17 +5837,31 @@ class Branches: async_table = LOOP.run( self._table.branches.create(name, from_ref, from_version) ) - return LanceTable.from_inner(async_table._inner) + return self._wrap(async_table) def checkout(self, name: str) -> "LanceTable": """Check out an existing branch and return a handle scoped to it.""" async_table = LOOP.run(self._table.branches.checkout(name)) - return LanceTable.from_inner(async_table._inner) + return self._wrap(async_table) def delete(self, name: str) -> None: """Delete a branch.""" LOOP.run(self._table.branches.delete(name)) + def _wrap(self, async_table: "AsyncTable") -> "LanceTable": + # Reuse the parent's connection + namespace context; from_inner would drop + # it and break identity/query routing for namespace-backed tables. + parent = self._parent + return LanceTable( + parent._conn, + async_table.name, + namespace_path=parent._namespace_path, + namespace_client=parent._namespace_client, + pushdown_operations=parent._pushdown_operations, + location=parent._location, + _async=async_table, + ) + class AsyncTags: """ diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b230969b9..931f8c759 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -942,6 +942,16 @@ def test_branches(tmp_path): assert "exp" not in table.branches.list() +def test_branches_preserve_namespace(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("t", [{"id": 1}], namespace_path=["ns1"]) + assert table.namespace == ["ns1"] + + branch = table.branches.create("exp") + assert branch.namespace == ["ns1"] + assert branch.id == table.id + + @pytest.mark.asyncio async def test_async_branches(tmp_path): db = await lancedb.connect_async(tmp_path)