diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 14f6dd25a..bc69c087a 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -544,6 +544,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, storage_options: Optional[Dict[str, str]] = None, index_cache_size: Optional[int] = None, + branch: Optional[str] = None, ) -> Table: if namespace_path is None: namespace_path = [] @@ -562,7 +563,7 @@ class LanceNamespaceDBConnection(DBConnection): raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}") raise - return LanceTable( + tbl = LanceTable( self, name, namespace_path=namespace_path, @@ -570,6 +571,9 @@ class LanceNamespaceDBConnection(DBConnection): pushdown_operations=self._namespace_client_pushdown_operations, _async=async_table, ) + if branch is not None: + return tbl.branches.checkout(branch) + return tbl @override def drop_table(self, name: str, namespace_path: Optional[List[str]] = None): diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 4421b057c..02fb5942f 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -383,6 +383,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, storage_options: Optional[Dict[str, str]] = None, index_cache_size: Optional[int] = None, + branch: Optional[str] = None, ) -> Table: """Open a Lance Table in the database. @@ -400,6 +401,9 @@ class RemoteDBConnection(DBConnection): """ from .table import RemoteTable + if branch is not None: + raise NotImplementedError("branching is not yet supported on remote tables") + if namespace_path is None: namespace_path = [] if index_cache_size is not None: