From d1ae41723c3c5603fecf29d84b2d0fb7c46f0417 Mon Sep 17 00:00:00 2001 From: Yue Zhang Date: Tue, 14 Jul 2026 02:23:12 -0400 Subject: [PATCH 1/3] perf(compaction): skip capturing row addresses when no index to remap --- rust/lance/src/dataset/optimize.rs | 83 +++++++++++++++++++++++++++--- 1 file changed, 77 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index bb51056eff5..0ae88ff7e1c 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1470,6 +1470,20 @@ async fn load_index_fragmaps(dataset: &Dataset) -> Result> { Ok(index_fragmaps) } +/// Whether the dataset has any index that compaction would need to remap. +/// +/// Only non-system indices are remapped per rewrite group; system indices +/// (fragment-reuse, mem-wal) are handled separately. When there are none, +/// compacting an address-style dataset has nothing to remap and +/// therefore does not need to capture row addresses at all. +async fn has_remappable_index(dataset: &Dataset) -> Result { + Ok(dataset + .load_indices() + .await? + .iter() + .any(|idx| !is_system_index(idx))) +} + pub async fn plan_compaction( dataset: &Dataset, options: &CompactionOptions, @@ -1570,8 +1584,11 @@ async fn rewrite_files( .iter() .map(|f| f.physical_rows.unwrap() as u64) .sum::(); - // If we aren't using stable row ids, then we need to remap indices. - let needs_remapping = !dataset.manifest.uses_stable_row_ids(); + // Capturing row addresses is only useful if something will consume them: an + // index to remap now, or deferred remap. With no index, skip it to avoid an + // extra scan column and a big treemap. + let capture_row_addrs = !dataset.manifest.uses_stable_row_ids() + && (options.defer_index_remap || has_remappable_index(dataset.as_ref()).await?); let mut new_fragments: Vec; let task_id = uuid::Uuid::new_v4(); log::info!( @@ -1597,7 +1614,7 @@ async fn rewrite_files( options.batch_size, options.io_buffer_size, true, - needs_remapping, + capture_row_addrs, ) .await?; row_ids_rx = rx_initial; @@ -1701,7 +1718,7 @@ async fn rewrite_files( )); } - if needs_remapping { + if capture_row_addrs { let (tx, rx) = std::sync::mpsc::channel(); let mut addrs = RoaringTreemap::new(); for frag in &fragments { @@ -1972,7 +1989,9 @@ pub async fn commit_compaction( } // If we aren't using stable row ids, then we need to remap indices. - let needs_remapping = !dataset.manifest.uses_stable_row_ids() && !options.defer_index_remap; + let has_address_style = completed_tasks.iter().any(|t| t.row_addrs.is_some()); + let needs_remapping = + !dataset.manifest.uses_stable_row_ids() && !options.defer_index_remap && has_address_style; // Determine the earliest version at which compaction tasks were planned/executed. // @@ -1996,7 +2015,6 @@ pub async fn commit_compaction( let mut completed_tasks = completed_tasks; // Single reserve_fragment_ids for all address-style tasks - let has_address_style = completed_tasks.iter().any(|t| t.row_addrs.is_some()); if has_address_style { let frags: Vec<&mut Fragment> = completed_tasks .iter_mut() @@ -2976,6 +2994,49 @@ mod tests { assert_eq!(dataset.manifest.uses_stable_row_ids(), use_stable_row_id,); } + #[rstest] + #[tokio::test] + async fn test_skip_rowid_capture_when_no_index(#[values(false, true)] has_index: bool) { + let test_dir = TempStrDir::default(); + + let data = sample_data(); + let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 9000))], data.schema()); + let write_params = WriteParams { + max_rows_per_file: 3000, // 3 fragments + data_storage_version: Some(LanceFileVersion::Legacy), + ..Default::default() + }; + let mut dataset = Dataset::write(reader, &test_dir, Some(write_params)) + .await + .unwrap(); + assert!(!dataset.manifest.uses_stable_row_ids()); + + if has_index { + dataset + .create_index( + &["a"], + IndexType::Scalar, + Some("a_idx".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + } + + let options = CompactionOptions { + target_rows_per_fragment: 9000, + ..Default::default() + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + + let result = rewrite_files(Cow::Borrowed(&dataset), plan.tasks()[0].clone(), &options) + .await + .unwrap(); + assert_eq!(result.row_addrs.is_some(), has_index); + } + #[tokio::test] async fn test_stable_row_indices() { // Validate behavior of indices after compaction with stable row ids. @@ -3236,6 +3297,16 @@ mod tests { ) .await .unwrap(); + dataset2 + .create_index( + &["i"], + IndexType::Scalar, + Some("scalar".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); // Verify the initial state - no fragment reuse index should exist let initial_indices = dataset.load_indices().await.unwrap(); From c6271aaa00dc2592f7bbda91e2a625460e40162e Mon Sep 17 00:00:00 2001 From: Yue Zhang Date: Wed, 15 Jul 2026 03:35:00 -0400 Subject: [PATCH 2/3] code review --- rust/lance/src/dataset/optimize.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index a9eb621b61e..f7c3c44deaa 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -3054,8 +3054,10 @@ mod tests { } #[rstest] + #[case::without_index(false)] + #[case::with_index(true)] #[tokio::test] - async fn test_skip_rowid_capture_when_no_index(#[values(false, true)] has_index: bool) { + async fn test_skip_rowid_capture_when_no_index(#[case] has_index: bool) { let test_dir = TempStrDir::default(); let data = sample_data(); @@ -3094,6 +3096,11 @@ mod tests { .await .unwrap(); assert_eq!(result.row_addrs.is_some(), has_index); + if let Some(row_addrs_bytes) = &result.row_addrs { + let row_addrs = + RoaringTreemap::deserialize_from(&mut Cursor::new(row_addrs_bytes)).unwrap(); + assert_eq!(row_addrs.len(), 9000); + } } #[tokio::test] From ef4e9704670e3e9621921d893080cf172772b83a Mon Sep 17 00:00:00 2001 From: Yue Zhang Date: Wed, 15 Jul 2026 04:05:05 -0400 Subject: [PATCH 3/3] code review --- rust/lance/src/dataset/optimize.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index f7c3c44deaa..e9cea451ff7 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1480,7 +1480,7 @@ async fn load_index_fragmaps(dataset: &Dataset) -> Result> { /// Only non-system indices are remapped per rewrite group; system indices /// (fragment-reuse, mem-wal) are handled separately. When there are none, /// compacting an address-style dataset has nothing to remap and -/// therefore does not need to capture row addresses at all. +/// therefore does not need to capture row addresses at all async fn has_remappable_index(dataset: &Dataset) -> Result { Ok(dataset .load_indices()