diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 975f3d9e..b1d27d95 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -48,7 +48,7 @@ class Table: async def version(self) -> int: ... async def checkout(self, version: int): ... async def checkout_latest(self): ... - async def restore(self): ... + async def restore(self, version: Optional[int] = None): ... async def list_indices(self) -> list[IndexConfig]: ... async def delete(self, filter: str): ... async def add_columns(self, columns: list[tuple[str, str]]) -> None: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 00dd9a0b..70747a21 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -87,6 +87,9 @@ class RemoteTable(Table): def checkout_latest(self): return LOOP.run(self._table.checkout_latest()) + def restore(self, version: Optional[int] = None): + return LOOP.run(self._table.restore(version)) + def list_indices(self) -> Iterable[IndexConfig]: """List all the indices on the table""" return LOOP.run(self._table.list_indices()) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 7fed0127..0beb4936 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1342,6 +1342,21 @@ class Table(ABC): It can also be used to undo a `[Self::checkout]` operation """ + @abstractmethod + def restore(self, version: Optional[int] = None): + """Restore a version of the table. This is an in-place operation. + + This creates a new version where the data is equivalent to the + specified previous version. Data is not copied (as of python-v0.2.1). + + Parameters + ---------- + version : int, default None + The version to restore. If unspecified then restores the currently + checked out version. If the currently checked out version is the + latest version then this is a no-op. + """ + @abstractmethod def list_versions(self) -> List[Dict[str, Any]]: """List all versions of the table""" @@ -3613,7 +3628,7 @@ class AsyncTable: """ await self._inner.checkout_latest() - async def restore(self): + async def restore(self, version: Optional[int] = None): """ Restore the table to the currently checked out version @@ -3626,7 +3641,7 @@ class AsyncTable: Once the operation concludes the table will no longer be in a checked out state and the read_consistency_interval, if any, will apply. """ - await self._inner.restore() + await self._inner.restore(version) async def optimize( self, diff --git a/python/src/table.rs b/python/src/table.rs index 08e83c86..9e3436f9 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -303,12 +303,16 @@ impl Table { }) } - pub fn restore(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (version=None))] + pub fn restore(self_: PyRef<'_, Self>, version: Option) -> PyResult> { let inner = self_.inner_ref()?.clone(); - future_into_py( - self_.py(), - async move { inner.restore().await.infer_error() }, - ) + + future_into_py(self_.py(), async move { + if let Some(version) = version { + inner.checkout(version).await.infer_error()?; + } + inner.restore().await.infer_error() + }) } pub fn query(&self) -> Query { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index e2bc0c8f..439bfb66 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -425,10 +425,16 @@ impl BaseTable for RemoteTable { Ok(()) } async fn restore(&self) -> Result<()> { - self.check_mutable().await?; - Err(Error::NotSupported { - message: "restore is not supported on LanceDB cloud.".into(), - }) + let mut request = self + .client + .post(&format!("/v1/table/{}/restore/", self.name)); + let version = self.current_version().await; + let body = serde_json::json!({ "version": version }); + request = request.json(&body); + + let (request_id, response) = self.client.send(request, true).await?; + self.check_table_response(&request_id, response).await?; + Ok(()) } async fn list_versions(&self) -> Result> {