diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 345d10f71db..69f5795fae6 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1508,6 +1508,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, @@ -1608,8 +1622,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!( @@ -1635,7 +1652,7 @@ async fn rewrite_files( options.batch_size, options.io_buffer_size, true, - needs_remapping, + capture_row_addrs, ) .await?; row_ids_rx = rx_initial; @@ -1739,7 +1756,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 { @@ -2010,7 +2027,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. // @@ -2034,7 +2053,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() @@ -3069,6 +3087,56 @@ mod tests { assert_eq!(dataset.manifest.uses_stable_row_ids(), use_stable_row_id,); } + #[rstest] + #[case::without_index(false)] + #[case::with_index(true)] + #[tokio::test] + async fn test_skip_rowid_capture_when_no_index(#[case] 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); + 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] async fn test_stable_row_indices() { // Validate behavior of indices after compaction with stable row ids. @@ -3333,6 +3401,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();