diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index a6c58bdd0..106b091e7 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -11,7 +11,8 @@ use std::{collections::HashSet, sync::Arc}; use arrow_schema::DataType; use lance::dataset::cleanup::RemovalStats; use lance::dataset::optimize::{ - CompactionMetrics, IndexRemapperOptions, compact_files, plan_compaction, + CompactionMetrics, CompactionPlan, CompactionPlanner, IndexRemapperOptions, + compact_files_with_planner, plan_compaction, }; use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; use lance_index::IndexType; @@ -26,6 +27,18 @@ pub use lance::dataset::optimize::CompactionOptions; const MAX_ARROW_FIXED_SIZE_LIST_CHILD_INDEX: u64 = u32::MAX as u64; +struct PrecomputedCompactionPlanner { + plan: CompactionPlan, +} + +#[async_trait::async_trait] +impl CompactionPlanner for PrecomputedCompactionPlanner { + async fn plan(&self, dataset: &lance::Dataset) -> lance::Result { + debug_assert_eq!(dataset.manifest().version, self.plan.read_version()); + Ok(self.plan.clone()) + } +} + fn sq_vector_dimension(data_type: &DataType) -> Option { match data_type { DataType::FixedSizeList(_, dimension) => Some(*dimension), @@ -63,15 +76,15 @@ async fn validate_sq_index_remapping( dataset: &lance::Dataset, options: &CompactionOptions, has_custom_remapper: bool, -) -> Result<()> { +) -> Result { + let plan = plan_compaction(dataset, options).await?; if options.defer_index_remap || has_custom_remapper || dataset.manifest().uses_stable_row_ids() { - return Ok(()); + return Ok(plan); } - let plan = plan_compaction(dataset, options).await?; if plan.tasks.is_empty() { - return Ok(()); + return Ok(plan); } let affected_fragments: HashSet = plan .tasks @@ -147,7 +160,7 @@ async fn validate_sq_index_remapping( } } - Ok(()) + Ok(plan) } /// Optimize the dataset. @@ -274,7 +287,7 @@ pub(crate) async fn cleanup_old_versions( /// This can be run after making several small appends to optimize the table /// for faster reads. /// -/// This calls into [lance::dataset::optimize::compact_files]. +/// This calls into [lance::dataset::optimize::compact_files_with_planner]. pub(crate) async fn compact_files_impl( table: &NativeTable, options: CompactionOptions, @@ -282,8 +295,9 @@ pub(crate) async fn compact_files_impl( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); - validate_sq_index_remapping(&dataset, &options, remap_options.is_some()).await?; - let metrics = compact_files(&mut dataset, options, remap_options).await?; + let plan = validate_sq_index_remapping(&dataset, &options, remap_options.is_some()).await?; + let planner = PrecomputedCompactionPlanner { plan }; + let metrics = compact_files_with_planner(&mut dataset, remap_options, &planner).await?; table.dataset.update(dataset); Ok(metrics) }