fix: address review comments on branch support

python/python/lancedb/namespace.py

@@ -978,12 +978,13 @@ class AsyncLanceNamespaceDBConnection:
         namespace_path: Optional[List[str]] = None,
         storage_options: Optional[Dict[str, str]] = None,
         index_cache_size: Optional[int] = None,
+        branch: Optional[str] = None,
     ) -> AsyncTable:
         """Open an existing table from the namespace."""
         if namespace_path is None:
             namespace_path = []
         try:
-            return await self._inner.open_table(
+            tbl = await self._inner.open_table(
                 name,
                 namespace_path=namespace_path,
                 storage_options=storage_options,
@@ -994,6 +995,9 @@ class AsyncLanceNamespaceDBConnection:
                 table_id = namespace_path + [name]
                 raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
             raise
+        if branch is not None:
+            return await tbl.branches.checkout(branch)
+        return tbl
 
     async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
         """Drop a table from the namespace."""

python/python/tests/test_table.py

@@ -982,6 +982,19 @@ def test_open_table_with_branch(tmp_path):
     assert db.open_table("t").count_rows() == 1
 
 
+@pytest.mark.asyncio
+async def test_async_namespace_open_table_with_branch(tmp_path):
+    db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
+    await db.create_namespace(["ns1"])
+    table = await db.create_table("t", [{"id": 1}], namespace_path=["ns1"])
+    branch = await table.branches.create("exp")
+    await branch.add([{"id": 2}])
+
+    # open_table(branch=...) on the async namespace connection must work
+    opened = await db.open_table("t", namespace_path=["ns1"], branch="exp")
+    assert await opened.count_rows() == 2
+
+
 def test_branch_to_lance_targets_branch(tmp_path):
     pytest.importorskip("lance")
     db = lancedb.connect(tmp_path)