fix: branch support on object stores, python branch ops, s3/ddb integration tests

python/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pylance"
 dynamic = ["version"]
-dependencies = ["pyarrow>=14", "numpy>=1.22", "lance-namespace>=0.8.0,<0.9"]
+dependencies = ["pyarrow>=14", "numpy>=1.22", "lance-namespace>=0.8.5,<0.9"]
 description = "python wrapper for Lance columnar format"
 authors = [{ name = "Lance Devs", email = "dev@lance.org" }]
 license = { file = "LICENSE" }

python/python/lance/dataset.py

@@ -950,6 +950,9 @@ def create_branch(
         ds._base_store_params = self._base_store_params
         ds._namespace_client = self._namespace_client
         ds._table_id = self._table_id
+        ds._namespace_client_managed_versioning = (
+            self._namespace_client_managed_versioning
+        )
         ds._default_scan_options = self._default_scan_options
         ds._read_params = self._read_params
         return ds
@@ -4579,6 +4582,7 @@ def commit_batch(
         ds._base_store_params = base_store_params
         ds._namespace_client = None
         ds._table_id = None
+        ds._namespace_client_managed_versioning = False
         ds._default_scan_options = None
         ds._read_params = None
         return BulkCommitResult(

python/python/lance/namespace.py

@@ -32,6 +32,8 @@
     CreateMaterializedViewResponse,
     CreateNamespaceRequest,
     CreateNamespaceResponse,
+    CreateTableBranchRequest,
+    CreateTableBranchResponse,
     CreateTableIndexRequest,
     CreateTableIndexResponse,
     CreateTableRequest,
@@ -42,6 +44,8 @@
     DeclareTableResponse,
     DeleteFromTableRequest,
     DeleteFromTableResponse,
+    DeleteTableBranchRequest,
+    DeleteTableBranchResponse,
     DeleteTableTagRequest,
     DeleteTableTagResponse,
     DeregisterTableRequest,
@@ -70,6 +74,8 @@
     LanceNamespace,
     ListNamespacesRequest,
     ListNamespacesResponse,
+    ListTableBranchesRequest,
+    ListTableBranchesResponse,
     ListTableIndicesRequest,
     ListTableIndicesResponse,
     ListTablesRequest,
@@ -850,6 +856,27 @@ def update_table_tag(
         response_dict = self._inner.update_table_tag(request.model_dump())
         return UpdateTableTagResponse.from_dict(response_dict)
 
+    def create_table_branch(
+        self, request: CreateTableBranchRequest
+    ) -> CreateTableBranchResponse:
+        """Create a new branch forked from a table version."""
+        response_dict = self._inner.create_table_branch(request.model_dump())
+        return CreateTableBranchResponse.from_dict(response_dict)
+
+    def list_table_branches(
+        self, request: ListTableBranchesRequest
+    ) -> ListTableBranchesResponse:
+        """List all branches of a table."""
+        response_dict = self._inner.list_table_branches(request.model_dump())
+        return ListTableBranchesResponse.from_dict(response_dict)
+
+    def delete_table_branch(
+        self, request: DeleteTableBranchRequest
+    ) -> DeleteTableBranchResponse:
+        """Delete a branch from a table."""
+        response_dict = self._inner.delete_table_branch(request.model_dump())
+        return DeleteTableBranchResponse.from_dict(response_dict)
+
     # Operation metrics methods
 
     def retrieve_ops_metrics(self) -> Dict[str, int]:
@@ -1420,6 +1447,27 @@ def update_table_tag(
         response_dict = self._inner.update_table_tag(request.model_dump())
         return UpdateTableTagResponse.from_dict(response_dict)
 
+    def create_table_branch(
+        self, request: CreateTableBranchRequest
+    ) -> CreateTableBranchResponse:
+        """Create a new branch forked from a table version."""
+        response_dict = self._inner.create_table_branch(request.model_dump())
+        return CreateTableBranchResponse.from_dict(response_dict)
+
+    def list_table_branches(
+        self, request: ListTableBranchesRequest
+    ) -> ListTableBranchesResponse:
+        """List all branches of a table."""
+        response_dict = self._inner.list_table_branches(request.model_dump())
+        return ListTableBranchesResponse.from_dict(response_dict)
+
+    def delete_table_branch(
+        self, request: DeleteTableBranchRequest
+    ) -> DeleteTableBranchResponse:
+        """Delete a branch from a table."""
+        response_dict = self._inner.delete_table_branch(request.model_dump())
+        return DeleteTableBranchResponse.from_dict(response_dict)
+
     # Operation metrics methods
 
     def retrieve_ops_metrics(self) -> Dict[str, int]:

python/python/tests/test_dataset.py

@@ -1742,6 +1742,7 @@ def test_commit_batch_append():
     result = lance.LanceDataset.commit_batch(dataset, [txn2, txn3])
     dataset = result["dataset"]
     assert dataset.version == 2
+    assert dataset.checkout_version(1).version == 1
     assert len(dataset.get_fragments()) == 3
     assert dataset.to_table() == pa.concat_tables([data1, data2, data3])
     merged_txn = result["merged"]
@@ -5538,6 +5539,8 @@ def test_branches(tmp_path: Path):
     branch1 = ds_main.create_branch("branch1")
     ds_main.branches.replace_metadata("branch1", {"description": "branch one"})
     assert branch1.version == 1
+    # The dataset returned by create_branch must be fully constructed
+    assert branch1.checkout_version(("main", None)).version == 1
     branch1_append = pa.Table.from_pydict({"a": [7, 8], "b": [9, 10]})
     branch1 = lance.write_dataset(branch1_append, branch1, mode="append")
     assert branch1.version == 2

python/python/tests/test_namespace_dir.py

@@ -29,6 +29,8 @@
     CountTableRowsRequest,
     CreateNamespaceRequest,
     CreateNamespaceResponse,
+    CreateTableBranchRequest,
+    CreateTableBranchResponse,
     CreateTableIndexRequest,
     CreateTableIndexResponse,
     CreateTableRequest,
@@ -37,6 +39,8 @@
     CreateTableVersionResponse,
     DeclareTableRequest,
     DeclareTableResponse,
+    DeleteTableBranchRequest,
+    DeleteTableBranchResponse,
     DeregisterTableRequest,
     DeregisterTableResponse,
     DescribeNamespaceRequest,
@@ -54,6 +58,8 @@
     InsertIntoTableResponse,
     ListNamespacesRequest,
     ListNamespacesResponse,
+    ListTableBranchesRequest,
+    ListTableBranchesResponse,
     ListTableIndicesRequest,
     ListTableIndicesResponse,
     ListTablesRequest,
@@ -71,6 +77,8 @@
     InvalidInputError,
     NamespaceNotEmptyError,
     NamespaceNotFoundError,
+    TableBranchAlreadyExistsError,
+    TableBranchNotFoundError,
     TableNotFoundError,
 )
 
@@ -151,6 +159,21 @@ def create_table_version(
     ) -> CreateTableVersionResponse:
         return self._inner.create_table_version(request)
 
+    def create_table_branch(
+        self, request: CreateTableBranchRequest
+    ) -> CreateTableBranchResponse:
+        return self._inner.create_table_branch(request)
+
+    def list_table_branches(
+        self, request: ListTableBranchesRequest
+    ) -> ListTableBranchesResponse:
+        return self._inner.list_table_branches(request)
+
+    def delete_table_branch(
+        self, request: DeleteTableBranchRequest
+    ) -> DeleteTableBranchResponse:
+        return self._inner.delete_table_branch(request)
+
     def create_table_index(
         self, request: CreateTableIndexRequest
     ) -> CreateTableIndexResponse:
@@ -564,6 +587,110 @@ def test_register_table_rejects_path_traversal(self, temp_ns_client):
         assert "Path traversal is not allowed" in str(exc_info.value)
 
 
+class TestTableBranchOperations:
+    """Branch CRUD through the python bindings - mirrors the Rust branch
+    CRUD tests."""
+
+    def test_branch_crud_round_trip(self, temp_ns_client):
+        create_ns_req = CreateNamespaceRequest(id=["workspace"])
+        temp_ns_client.create_namespace(create_ns_req)
+        ipc_data = table_to_ipc_bytes(create_test_data())
+        table_id = ["workspace", "branched_table"]
+        temp_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data)
+
+        temp_ns_client.create_table_branch(
+            CreateTableBranchRequest(id=table_id, name="dev")
+        )
+        listed = temp_ns_client.list_table_branches(
+            ListTableBranchesRequest(id=table_id)
+        )
+        assert "dev" in listed.branches
+        assert listed.branches["dev"].parent_version == 1
+
+        # Duplicate creation and deleting a missing branch surface the typed
+        # branch errors (codes 23 and 22), not InternalError.
+        temp_ns_client.create_table_branch(
+            CreateTableBranchRequest(id=table_id, name="dev2")
+        )
+        with pytest.raises(TableBranchAlreadyExistsError):
+            temp_ns_client.create_table_branch(
+                CreateTableBranchRequest(id=table_id, name="dev2")
+            )
+
+        temp_ns_client.delete_table_branch(
+            DeleteTableBranchRequest(id=table_id, name="dev")
+        )
+        listed = temp_ns_client.list_table_branches(
+            ListTableBranchesRequest(id=table_id)
+        )
+        assert "dev" not in listed.branches
+        with pytest.raises(TableBranchNotFoundError):
+            temp_ns_client.delete_table_branch(
+                DeleteTableBranchRequest(id=table_id, name="dev")
+            )
+
+    def test_create_branch_from_other_branch(self, temp_ns_client):
+        """Forking from a non-main source branch records the right parent."""
+        create_ns_req = CreateNamespaceRequest(id=["workspace"])
+        temp_ns_client.create_namespace(create_ns_req)
+        ipc_data = table_to_ipc_bytes(create_test_data())
+        table_id = ["workspace", "fork_table"]
+        temp_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data)
+
+        temp_ns_client.create_table_branch(
+            CreateTableBranchRequest(id=table_id, name="dev")
+        )
+        temp_ns_client.create_table_branch(
+            CreateTableBranchRequest(id=table_id, name="child", from_branch="dev")
+        )
+        listed = temp_ns_client.list_table_branches(
+            ListTableBranchesRequest(id=table_id)
+        )
+        assert listed.branches["child"].parent_branch == "dev"
+
+
+class _ForeignCodeError(Exception):
+    """Not a LanceNamespaceError, but carries the same integer code as
+    TABLE_NOT_FOUND."""
+
+    code = 4
+
+
+class _RaisingNamespace(LanceNamespace):
+    """A namespace whose describe_table raises the configured exception."""
+
+    def __init__(self, exc: Exception):
+        self._exc = exc
+
+    def namespace_id(self) -> str:
+        return "raising"
+
+    def describe_table(self, request: DescribeTableRequest) -> DescribeTableResponse:
+        raise self._exc
+
+
+class TestPythonNamespaceErrorMapping:
+    """The Rust adapter must trust the `code` attribute only on the
+    lance_namespace exception hierarchy."""
+
+    def test_namespace_error_identity_preserved(self):
+        ns = _RaisingNamespace(TableNotFoundError("no such table"))
+        with pytest.raises(TableNotFoundError, match="no such table"):
+            lance.dataset(namespace_client=ns, table_id=["t"])
+
+        # Branch error codes (22/23) survive the round trip too.
+        ns = _RaisingNamespace(TableBranchNotFoundError("no such branch"))
+        with pytest.raises(TableBranchNotFoundError, match="no such branch"):
+            lance.dataset(namespace_client=ns, table_id=["t"])
+
+    def test_foreign_code_attribute_not_trusted(self):
+        # The foreign exception must surface as itself, not be reinterpreted
+        # as a namespace error via its `code` attribute.
+        ns = _RaisingNamespace(_ForeignCodeError("boom"))
+        with pytest.raises(_ForeignCodeError, match="boom"):
+            lance.dataset(namespace_client=ns, table_id=["t"])
+
+
 class TestChildNamespaceOperations:
     """Tests for operations in child namespaces - mirrors Rust tests."""