diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 7d10f8b59..84340c72c 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3742,6 +3742,21 @@ describe("when optimizing a dataset", () => { expect(stats.prune.oldVersionsRemoved).toBe(2); }); + it("retains tagged versions older than the cleanup cutoff", async () => { + await (await table.tags()).create("keep", 1); + await table.add([{ id: 3 }]); + + const cutoff = await nextMillisecond(); + const stats = await table.optimize({ cleanupOlderThan: cutoff }); + // Versions 2 and 3 are pruned, tagged version 1 is retained + expect(stats.prune.oldVersionsRemoved).toBe(2); + + await table.checkout("keep"); + expect(await table.countRows()).toBe(1); + await table.checkoutLatest(); + expect(await table.countRows()).toBe(3); + }); + it("delete unverified", async () => { const version = await table.version(); const versionFile = `${tmpDir.name}/${table.name}.lance/_versions/${String( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index e23b478f5..b55abfa7e 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4792,6 +4792,57 @@ def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection): assert all("concurrent" in str(w.message) for w in seen) +def test_optimize_retains_tagged_versions(tmp_db: DBConnection): + # A tagged version older than the cleanup cutoff must be retained, not make + # optimize() fail part-way (after compaction, before the index step). + table = tmp_db.create_table( + "test", data=pa.table({"id": list(range(100)), "text": ["a"] * 100}) + ) + table.create_index("id", config=BTree()) + table.tags.create("release-1", table.version) + tagged_version = table.version + for i in range(3): + table.add(pa.table({"id": [1000 + i], "text": ["new"]})) + assert table.stats()["fragment_stats"]["num_fragments"] == 4 + assert table.index_stats("id_idx").num_unindexed_rows == 3 + sleep(1.1) + + table.optimize(cleanup_older_than=timedelta(seconds=1)) + + # The index step ran after the prune step + assert table.index_stats("id_idx").num_unindexed_rows == 0 + # It stays optimizable + table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert table.stats()["fragment_stats"]["num_fragments"] == 1 + + # Untagged old versions were pruned, the tagged one was retained + versions = [v["version"] for v in table.list_versions()] + assert tagged_version in versions + for added in range(tagged_version + 1, tagged_version + 4): + assert added not in versions + assert len(table) == 103 + + table.checkout("release-1") + assert len(table) == 100 + assert table.to_arrow().num_rows == 100 + + +@pytest.mark.asyncio +async def test_optimize_retains_tagged_versions_async(tmp_db_async: AsyncConnection): + table = await tmp_db_async.create_table("test", data=[{"x": [1]}]) + await table.tags.create("keep", await table.version()) + await table.add(data=[{"x": [2]}]) + await table.add(data=[{"x": [3]}]) + sleep(1.1) + + stats = await table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert stats.prune.old_versions_removed >= 1 + + await table.checkout("keep") + assert await table.count_rows() == 1 + assert await table.query().to_arrow() == pa.table({"x": [[1]]}) + + def test_replace_field_metadata(tmp_path): db = lancedb.connect(tmp_path) table = db.create_table("my_table", data=[{"x": 0}]) diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index 2a502b059..ebe03359c 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -69,7 +69,11 @@ pub enum OptimizeAction { /// **WARNING**: This should only be set to true if you can guarantee that no other process is /// currently working on this dataset. Otherwise the dataset could be put into a corrupted state. delete_unverified: Option, + /// Tagged versions are never removed, regardless of their age. + /// /// If true, an error will be returned if there are any old versions that are still tagged. + /// If false or `None` (the default), tagged versions are silently retained and all other + /// old versions are removed. error_if_tagged_old_versions: Option, }, /// Optimize the indices @@ -125,6 +129,12 @@ pub(crate) async fn optimize_indices(table: &NativeTable, options: &OptimizeOpti /// no other process is currently working on this dataset. Otherwise the /// dataset could be put into a corrupted state. /// +/// * `error_if_tagged_old_versions` - Tagged versions are always retained. If +/// this is true, an error is returned when a tagged version is older than the +/// cutoff. Defaults to false (the tagged versions are retained silently), +/// unlike Lance's own default, so that `optimize` does not fail (and skip the +/// steps after pruning) once a tagged version ages past the cutoff. +/// /// This calls into [lance::dataset::Dataset::cleanup_old_versions] and /// returns the result. pub(crate) async fn cleanup_old_versions( @@ -136,7 +146,11 @@ pub(crate) async fn cleanup_old_versions( table.dataset.ensure_mutable()?; let dataset = table.dataset.get().await?; let stats = dataset - .cleanup_old_versions(older_than, delete_unverified, error_if_tagged_old_versions) + .cleanup_old_versions( + older_than, + delete_unverified, + Some(error_if_tagged_old_versions.unwrap_or(false)), + ) .await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -161,9 +175,8 @@ pub(crate) async fn cleanup_old_versions_before( if let Some(delete_unverified) = delete_unverified { policy = policy.delete_unverified(delete_unverified); } - if let Some(error_if_tagged_old_versions) = error_if_tagged_old_versions { - policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions); - } + // Tagged versions are retained rather than raising an error unless asked. + policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions.unwrap_or(false)); let stats = dataset.cleanup_with_policy(policy.build()).await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -511,6 +524,90 @@ mod tests { assert_eq!(all_values, expected); } + #[tokio::test] + async fn test_optimize_prune_retains_tagged_versions() { + let tmpdir = tempfile::tempdir().unwrap(); + let conn = connect(tmpdir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let make_batch = |start: i32| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(start..start + 10))], + ) + .unwrap() + }; + + let table = conn + .create_table("test_prune_tagged", make_batch(0)) + .execute() + .await + .unwrap(); + // Tag the first version, then create more versions after it. + table.tags().await.unwrap().create("keep", 1).await.unwrap(); + for i in 1..4 { + table.add(make_batch(i * 10)).execute().await.unwrap(); + } + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // Explicitly asking for an error still errors, and leaves everything in place. + let err = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: Some(true), + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("tagged version"), "{err}"); + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // By default, tagged versions are retained and the other old versions are removed. + let stats = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: None, + }) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 2); + let mut versions = table + .list_versions() + .await + .unwrap() + .into_iter() + .map(|v| v.version) + .collect::>(); + versions.sort(); + assert_eq!(versions, vec![1, 4]); + + // Same through the absolute-cutoff entry point used by the Node binding. + table.add(make_batch(40)).execute().await.unwrap(); + let stats = table + .optimize_prune_before(chrono::Utc::now(), Some(true), None) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 1); + + // The tagged version is still readable, including its data files. + assert_eq!(table.count_rows(None).await.unwrap(), 50); + table.checkout_tag("keep").await.unwrap(); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 10); + } + #[tokio::test] async fn test_optimize_index() { let conn = connect("memory://").execute().await.unwrap();