mirror of
https://github.com/lancedb/lancedb.git
synced 2026-07-13 16:00:43 +00:00
feat: support checking out a version on a branch (#3504)
### Description Stacked on #3490. Adds an optional version to branch checkout across the Rust core and the Python and TypeScript SDKs, so you can open a specific version on a branch ("version V of branch B"), not just the branch's latest version Rust ```rust // Open version 3 of branch "exp" (a read-only view): check out from an // existing table, or open it directly from the connection. let exp_v3 = table.checkout_branch("exp", Some(3)).await?; let exp_v3 = db.open_table("items").branch("exp").version(3).execute().await?; // checkout_latest re-attaches to the branch's writable HEAD. exp_v3.checkout_latest().await?; // With no branch, a version opens main at that version. let main_v3 = db.open_table("items").version(3).execute().await?; ``` Python ```python # Open version 3 of branch "exp" (a read-only view): check out from an # existing table, or open it directly from the connection. branch_v3 = await table.branches.checkout("exp", version=3) branch_v3 = await db.open_table("items", branch="exp", version=3) # checkout_latest re-attaches to the branch's writable HEAD. await branch_v3.checkout_latest() # With no branch, a version opens main at that version. main_v3 = await db.open_table("items", version=3) ``` TypeScript ```typescript // Open version 3 of branch "exp" (a read-only view): check out from an // existing table, or open it directly from the connection. const branchV3 = await (await table.branches()).checkout("exp", 3); const opened = await db.openTable("items", undefined, { branch: "exp", version: 3 }); // checkoutLatest re-attaches to the branch's writable HEAD. await branchV3.checkoutLatest(); // With no branch, a version opens main at that version. const mainV3 = await db.openTable("items", undefined, { version: 3 }); ``` ### Testing - Added unit tests (Rust, Python sync + async, TypeScript): branch-scoped resolution at a version number shared with `main` and with another branch, read-only enforcement on a pinned handle, `checkout_latest` recovery to the branch's HEAD, fork-point reads, and the nonexistent-version/branch error paths. - Ran smoke tests against the Python and TypeScript SDKs on local machine.
This commit is contained in:
Brendan Clement
GitHub
@@ -9,6 +9,7 @@ use std::sync::Arc;
|
||||
use arrow_array::RecordBatch;
|
||||
use arrow_schema::SchemaRef;
|
||||
use lance::dataset::ReadParams;
|
||||
use lance::dataset::refs::MAIN_BRANCH;
|
||||
use lance_namespace::models::{
|
||||
CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest,
|
||||
DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest,
|
||||
@@ -120,6 +121,7 @@ pub struct OpenTableBuilder {
|
||||
request: OpenTableRequest,
|
||||
embedding_registry: Arc<dyn EmbeddingRegistry>,
|
||||
branch: Option<String>,
|
||||
version: Option<u64>,
|
||||
}
|
||||
|
||||
impl OpenTableBuilder {
|
||||
@@ -141,6 +143,7 @@ impl OpenTableBuilder {
|
||||
},
|
||||
embedding_registry,
|
||||
branch: None,
|
||||
version: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -269,13 +272,39 @@ impl OpenTableBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Open the table pinned to a specific version, producing a read-only "view".
|
||||
///
|
||||
/// Composes with [`Self::branch`]: when a branch is also set, this opens that
|
||||
/// branch at the given version; otherwise it opens `main` at that version.
|
||||
/// The returned table is a detached head, so operations that modify the table
|
||||
/// will fail until [`Table::checkout_latest`] is called.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::Connection;
|
||||
/// # async fn f(conn: &Connection) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let table = conn.open_table("t").branch("exp").version(3).execute().await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn version(mut self, version: u64) -> Self {
|
||||
self.version = Some(version);
|
||||
self
|
||||
}
|
||||
|
||||
/// Open the table
|
||||
pub async fn execute(self) -> Result<Table> {
|
||||
let table = self.parent.open_table(self.request).await?;
|
||||
let table = Table::new_with_embedding_registry(table, self.parent, self.embedding_registry);
|
||||
match self.branch {
|
||||
Some(branch) => table.checkout_branch(&branch).await,
|
||||
None => Ok(table),
|
||||
// "main" is the default branch, so treat it as no branch.
|
||||
let branch = self.branch.filter(|b| b.as_str() != MAIN_BRANCH);
|
||||
match branch {
|
||||
Some(branch) => table.checkout_branch(&branch, self.version).await,
|
||||
None => {
|
||||
if let Some(version) = self.version {
|
||||
table.checkout(version).await?;
|
||||
}
|
||||
Ok(table)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -983,6 +983,49 @@ mod tests {
|
||||
assert_eq!(table.name(), "table1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_table_branch_and_version() {
|
||||
// Remote supports version time-travel but not branches. A version-only
|
||||
// open (or one on the default "main" branch) must succeed; a non-main
|
||||
// branch must be rejected, with or without a version.
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/t/describe/");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"table": "t", "version": 2, "schema": {"fields": [
|
||||
{"name": "a", "type": { "type": "int32" }, "nullable": false}
|
||||
]}}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
// version-only: allowed (open + checkout(version) both round-trip)
|
||||
conn.open_table("t").version(2).execute().await.unwrap();
|
||||
|
||||
// "main" is the default branch, so it counts as no branch
|
||||
conn.open_table("t")
|
||||
.branch("main")
|
||||
.version(2)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// a non-main branch is rejected, with or without a version
|
||||
assert!(matches!(
|
||||
conn.open_table("t").branch("exp").execute().await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
conn.open_table("t")
|
||||
.branch("exp")
|
||||
.version(2)
|
||||
.execute()
|
||||
.await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_table_not_found() {
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
|
||||
@@ -633,6 +633,23 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Check out an existing branch and return a handle scoped to it.
|
||||
async fn checkout_branch(&self, name: &str) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Check out an existing branch at an optional version, returning a handle.
|
||||
///
|
||||
/// `None` tracks the branch's latest; `Some(v)` pins it to that version
|
||||
/// (read-only). The default implementation composes [`Self::checkout_branch`]
|
||||
/// and [`Self::checkout`]; implementations may override it to resolve the
|
||||
/// `(branch, version)` coordinate in a single manifest read.
|
||||
async fn checkout_branch_version(
|
||||
&self,
|
||||
name: &str,
|
||||
version: Option<u64>,
|
||||
) -> Result<Arc<dyn BaseTable>> {
|
||||
let branch = self.checkout_branch(name).await?;
|
||||
if let Some(version) = version {
|
||||
branch.checkout(version).await?;
|
||||
}
|
||||
Ok(branch)
|
||||
}
|
||||
/// List the branches of the table.
|
||||
async fn list_branches(&self) -> Result<HashMap<String, BranchContents>>;
|
||||
/// Delete a branch.
|
||||
@@ -1654,8 +1671,20 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Check out an existing branch and return a handle scoped to it.
|
||||
pub async fn checkout_branch(&self, name: &str) -> Result<Self> {
|
||||
let inner = self.inner.checkout_branch(name).await?;
|
||||
///
|
||||
/// With `version` set, the returned handle is pinned to that version of the
|
||||
/// branch: a read-only, detached view (as with [`Self::checkout`]). With
|
||||
/// `version` as `None` it tracks the branch's latest and stays writable.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::Table;
|
||||
/// # async fn f(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let exp_at_v3 = table.checkout_branch("exp", Some(3)).await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn checkout_branch(&self, name: &str, version: Option<u64>) -> Result<Self> {
|
||||
let inner = self.inner.checkout_branch_version(name, version).await?;
|
||||
Ok(Self {
|
||||
inner,
|
||||
database: self.database.clone(),
|
||||
@@ -2757,6 +2786,29 @@ impl BaseTable for NativeTable {
|
||||
Ok(Arc::new(self.with_dataset(dataset)))
|
||||
}
|
||||
|
||||
async fn checkout_branch_version(
|
||||
&self,
|
||||
name: &str,
|
||||
version: Option<u64>,
|
||||
) -> Result<Arc<dyn BaseTable>> {
|
||||
let Some(version) = version else {
|
||||
return self.checkout_branch(name).await;
|
||||
};
|
||||
Self::validate_branch_name(name, "branch name")?;
|
||||
// Resolve (branch, version) in a single manifest read.
|
||||
let branch_ds = self
|
||||
.dataset
|
||||
.get()
|
||||
.await?
|
||||
.checkout_version((name, version))
|
||||
.await?;
|
||||
let dataset = dataset::DatasetConsistencyWrapper::new_time_travel(
|
||||
branch_ds,
|
||||
self.read_consistency_interval,
|
||||
);
|
||||
Ok(Arc::new(self.with_dataset(dataset)))
|
||||
}
|
||||
|
||||
async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
|
||||
Ok(self.dataset.get().await?.list_branches().await?)
|
||||
}
|
||||
@@ -3530,7 +3582,7 @@ mod tests {
|
||||
assert!(branches.contains_key("exp"));
|
||||
|
||||
// checking out the branch from the main handle sees the branch's latest data
|
||||
let checked_out = table.checkout_branch("exp").await.unwrap();
|
||||
let checked_out = table.checkout_branch("exp", None).await.unwrap();
|
||||
assert_eq!(checked_out.current_branch().as_deref(), Some("exp"));
|
||||
assert_eq!(checked_out.count_rows(None).await.unwrap(), 2);
|
||||
|
||||
@@ -3550,6 +3602,186 @@ mod tests {
|
||||
assert!(!branches.contains_key("exp"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_branch_version_checkout() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let uri = tmp_dir.path().to_str().unwrap();
|
||||
|
||||
let conn = ConnectBuilder::new(uri)
|
||||
.read_consistency_interval(Duration::from_secs(0))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// main: a single fork-point row (i = 0)
|
||||
let table = conn
|
||||
.create_table("my_table", sample_rows(vec![0]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let fork_point = table.version().await.unwrap();
|
||||
|
||||
// Fork "exp", then advance exp AND main independently past the fork so
|
||||
// they diverge while sharing version numbers.
|
||||
let branch = table.create_branch("exp", fork_point).await.unwrap();
|
||||
let exp_fork = branch.version().await.unwrap(); // exp's shallow-clone version
|
||||
branch.add(sample_rows(vec![1])).execute().await.unwrap(); // exp: {0, 1}
|
||||
let exp_v2 = branch.version().await.unwrap();
|
||||
branch.add(sample_rows(vec![2])).execute().await.unwrap(); // exp HEAD: {0, 1, 2}
|
||||
|
||||
// main's own commit reaches the SAME version number with different data
|
||||
table
|
||||
.add(sample_rows(vec![100, 101, 102]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap(); // main HEAD: {0, 100, 101, 102}
|
||||
let main_v2 = table.version().await.unwrap();
|
||||
assert_eq!(
|
||||
exp_v2, main_v2,
|
||||
"branch and main must share the version number for this test to mean anything"
|
||||
);
|
||||
|
||||
// Open exp at the shared version. The data must be exp's, not main's:
|
||||
// count alone cannot prove this (main@v2 differs), so assert provenance
|
||||
// by content.
|
||||
let pinned = conn
|
||||
.open_table("my_table")
|
||||
.branch("exp")
|
||||
.version(exp_v2)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
|
||||
// isolated from exp's HEAD (3 rows) and from main@v2 (4 rows)
|
||||
assert_eq!(pinned.count_rows(None).await.unwrap(), 2);
|
||||
// exp's post-fork row is visible; main's divergent rows are not
|
||||
assert_eq!(
|
||||
pinned.count_rows(Some("i = 1".to_string())).await.unwrap(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
pinned
|
||||
.count_rows(Some("i = 100".to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
|
||||
// the same coordinate is reachable directly via checkout_branch(name, version)
|
||||
let pinned_direct = table.checkout_branch("exp", Some(exp_v2)).await.unwrap();
|
||||
assert_eq!(pinned_direct.current_branch().as_deref(), Some("exp"));
|
||||
assert_eq!(pinned_direct.count_rows(None).await.unwrap(), 2);
|
||||
|
||||
// the HEADs are unaffected
|
||||
let head = conn
|
||||
.open_table("my_table")
|
||||
.branch("exp")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(head.count_rows(None).await.unwrap(), 3);
|
||||
assert_eq!(table.count_rows(None).await.unwrap(), 4);
|
||||
|
||||
// a pinned version is a detached head: writes are rejected
|
||||
assert!(pinned.add(sample_rows(vec![9])).execute().await.is_err());
|
||||
|
||||
// version-only (no branch) time-travels main itself: its fork-point
|
||||
// version holds only main's first row, and the shared version number
|
||||
// resolves to main's data, not the branch's ("opens main at the version")
|
||||
let old_main = conn
|
||||
.open_table("my_table")
|
||||
.version(fork_point)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(old_main.current_branch(), None);
|
||||
assert_eq!(old_main.count_rows(None).await.unwrap(), 1);
|
||||
let shared_on_main = conn
|
||||
.open_table("my_table")
|
||||
.version(exp_v2)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(shared_on_main.current_branch(), None);
|
||||
assert_eq!(shared_on_main.count_rows(None).await.unwrap(), 4);
|
||||
|
||||
// a nonexistent version is rejected
|
||||
assert!(
|
||||
conn.open_table("my_table")
|
||||
.version(9999)
|
||||
.execute()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
|
||||
// a nonexistent version on a branch is rejected too: this resolves on
|
||||
// the branch's path, a distinct miss from the main lookup above
|
||||
assert!(
|
||||
conn.open_table("my_table")
|
||||
.branch("exp")
|
||||
.version(9999)
|
||||
.execute()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
|
||||
// opening the branch at its fork point (the shallow-clone manifest)
|
||||
// shows just the cloned state: main's fork-point row
|
||||
let exp_at_fork = conn
|
||||
.open_table("my_table")
|
||||
.branch("exp")
|
||||
.version(exp_fork)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(exp_at_fork.current_branch().as_deref(), Some("exp"));
|
||||
assert_eq!(exp_at_fork.count_rows(None).await.unwrap(), 1);
|
||||
|
||||
// checkout_latest re-attaches the pinned handle to the BRANCH's HEAD
|
||||
// (writable again), not main's HEAD, and not staying pinned
|
||||
pinned.checkout_latest().await.unwrap();
|
||||
assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
|
||||
assert_eq!(pinned.count_rows(None).await.unwrap(), 3); // exp HEAD, not main's 4
|
||||
pinned.add(sample_rows(vec![3])).execute().await.unwrap();
|
||||
assert_eq!(pinned.count_rows(None).await.unwrap(), 4); // writable again
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_branch_version_two_branches() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let uri = tmp_dir.path().to_str().unwrap();
|
||||
let conn = ConnectBuilder::new(uri)
|
||||
.read_consistency_interval(Duration::from_secs(0))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let table = conn
|
||||
.create_table("my_table", sample_rows(vec![0]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let fork_point = table.version().await.unwrap();
|
||||
|
||||
// two branches off the same point, each advanced once so they reach the
|
||||
// SAME version number with divergent data
|
||||
let exp1 = table.create_branch("exp1", fork_point).await.unwrap();
|
||||
let exp2 = table.create_branch("exp2", fork_point).await.unwrap();
|
||||
exp1.add(sample_rows(vec![10])).execute().await.unwrap();
|
||||
exp2.add(sample_rows(vec![20])).execute().await.unwrap();
|
||||
let v1 = exp1.version().await.unwrap();
|
||||
let v2 = exp2.version().await.unwrap();
|
||||
assert_eq!(v1, v2, "both branches must reach the same version number");
|
||||
|
||||
// that shared version number resolves to each branch's own data
|
||||
let at1 = table.checkout_branch("exp1", Some(v1)).await.unwrap();
|
||||
assert_eq!(at1.count_rows(Some("i = 10".to_string())).await.unwrap(), 1);
|
||||
assert_eq!(at1.count_rows(Some("i = 20".to_string())).await.unwrap(), 0);
|
||||
let at2 = table.checkout_branch("exp2", Some(v2)).await.unwrap();
|
||||
assert_eq!(at2.count_rows(Some("i = 20".to_string())).await.unwrap(), 1);
|
||||
assert_eq!(at2.count_rows(Some("i = 10".to_string())).await.unwrap(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_branch_name_validation() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
@@ -3567,7 +3799,7 @@ mod tests {
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
table.checkout_branch("").await,
|
||||
table.checkout_branch("", None).await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
@@ -4005,6 +4237,19 @@ mod tests {
|
||||
Box::new(RecordBatchIterator::new(vec![batch], schema))
|
||||
}
|
||||
|
||||
/// A single-batch reader holding the given `i` (Int32) values. Lets a test
|
||||
/// write distinguishable rows so it can assert data provenance, not row count.
|
||||
fn sample_rows(values: Vec<i32>) -> Box<dyn arrow_array::RecordBatchReader + Send> {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(values))],
|
||||
)
|
||||
.unwrap();
|
||||
let schema = batch.schema().clone();
|
||||
|
||||
Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_scalar_index() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
@@ -76,6 +76,23 @@ impl DatasetConsistencyWrapper {
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new wrapper pinned to the dataset's current version.
|
||||
///
|
||||
/// `dataset` must already be checked out at the desired version; this pins
|
||||
/// to `dataset.version()` without re-resolving. The wrapper is read-only
|
||||
/// (time-travel) until [`as_latest`](Self::as_latest) re-attaches it to the
|
||||
/// latest version.
|
||||
pub fn new_time_travel(dataset: Dataset, read_consistency_interval: Option<Duration>) -> Self {
|
||||
let version = dataset.version().version;
|
||||
let wrapper = Self::new_latest(dataset, read_consistency_interval);
|
||||
wrapper
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pinned_version = Some(version);
|
||||
wrapper
|
||||
}
|
||||
|
||||
/// The MemWAL `ShardWriter` cache co-located with this dataset.
|
||||
pub(crate) fn shard_writer(&self) -> &Arc<ShardWriterCache> {
|
||||
&self.shard_writer
|
||||
|
||||
Reference in New Issue
Block a user