diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index 7cac7815125..bcbe8bbbbd2 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -17,7 +17,7 @@ use arrow_schema::DataType; use async_trait::async_trait; use lance_core::{Error, Result}; use lance_encoding::version::LanceFileVersion; -use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; +use lance_index::is_system_index; use lance_index::pb::VectorIndexDetails; use lance_index::scalar::lance_format::LanceIndexStore; use lance_table::format::IndexMetadata; @@ -29,19 +29,39 @@ use super::optimize::{IndexRemapper, IndexRemapperOptions}; #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct DatasetIndexRemapperOptions {} +/// Loads index metadata when compaction has at least one index to remap. +/// +/// Returns all index metadata, including system indices, so the remapper uses a +/// consistent snapshot. Returns `None` when there are no non-system indices. +pub(crate) async fn load_indices_for_remapping( + dataset: &Dataset, +) -> Result>>> { + if dataset.manifest.index_section.is_none() { + return Ok(None); + } + + let indices = dataset.load_indices().await?; + let has_remappable_index = indices.iter().any(|index| !is_system_index(index)); + Ok(has_remappable_index.then_some(indices)) +} + +#[async_trait] impl IndexRemapperOptions for DatasetIndexRemapperOptions { - fn create_remapper( - &self, - dataset: &Dataset, - ) -> crate::Result> { - Ok(Box::new(DatasetIndexRemapper { + async fn create_remapper(&self, dataset: &Dataset) -> Result>> { + let Some(indices) = load_indices_for_remapping(dataset).await? else { + return Ok(None); + }; + + Ok(Some(Box::new(DatasetIndexRemapper { dataset: Arc::new(dataset.clone()), - })) + indices, + }))) } } struct DatasetIndexRemapper { dataset: Arc, + indices: Arc>, } impl DatasetIndexRemapper { @@ -62,10 +82,9 @@ impl IndexRemapper for DatasetIndexRemapper { affected_fragment_ids: &[u64], ) -> Result> { let affected_frag_ids = HashSet::::from_iter(affected_fragment_ids.iter().copied()); - let indices = self.dataset.load_indices().await?; - let mut remapped = Vec::with_capacity(indices.len()); - for index in indices.iter() { - let needs_remapped = index.name != FRAG_REUSE_INDEX_NAME + let mut remapped = Vec::with_capacity(self.indices.len()); + for index in self.indices.iter() { + let needs_remapped = !is_system_index(index) && match &index.fragment_bitmap { None => true, Some(fragment_bitmap) => fragment_bitmap @@ -178,13 +197,55 @@ impl LanceIndexStoreExt for LanceIndexStore { mod tests { use super::*; use crate::dataset::WriteParams; + use crate::dataset::transaction::{Operation, Transaction}; use crate::index::DatasetIndexExt; + use crate::index::frag_reuse::build_frag_reuse_index_metadata; use crate::index::vector::VectorIndexParams; use lance_datagen::{BatchCount, RowCount, array}; use lance_index::IndexType; + use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragReuseIndexDetails}; use lance_linalg::distance::MetricType; use uuid::Uuid; + #[tokio::test] + async fn test_remapper_not_created_without_remappable_indices() { + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_reader_rows(RowCount::from(1), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + let options = DatasetIndexRemapperOptions::default(); + + assert!(options.create_remapper(&dataset).await.unwrap().is_none()); + + let frag_reuse_index = build_frag_reuse_index_metadata( + &dataset, + None, + FragReuseIndexDetails { + versions: Vec::new(), + }, + Default::default(), + ) + .await + .unwrap(); + let transaction = Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![frag_reuse_index], + removed_indices: Vec::new(), + }, + None, + ); + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].name, FRAG_REUSE_INDEX_NAME); + assert!(options.create_remapper(&dataset).await.unwrap().is_none()); + } + #[tokio::test] async fn test_remapper_only_touches_segments_with_affected_fragments() { let test_dir = tempfile::tempdir().unwrap(); @@ -296,7 +357,9 @@ mod tests { let remapper = DatasetIndexRemapperOptions::default() .create_remapper(&dataset) - .unwrap(); + .await + .unwrap() + .expect("vector index should require a remapper"); let remapped = remapper .remap_indices(RowAddrRemap::empty(), &[target_fragments[0].id() as u64]) .await diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 345d10f71db..17a9acaaaf9 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -88,7 +88,7 @@ use std::ops::{AddAssign, Range}; use std::sync::Arc; use super::fragment::FileFragment; -use super::index::DatasetIndexRemapperOptions; +use super::index::{DatasetIndexRemapperOptions, load_indices_for_remapping}; use super::rowids::load_row_id_sequences; use super::transaction::{ Operation, RewriteGroup, RewrittenIndex, Transaction, TransactionBuilder, @@ -1608,8 +1608,13 @@ 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 a deferred remap through the FRI. + let capture_row_addrs = !dataset.manifest.uses_stable_row_ids() + && (options.defer_index_remap + || load_indices_for_remapping(dataset.as_ref()) + .await? + .is_some()); let mut new_fragments: Vec; let task_id = uuid::Uuid::new_v4(); log::info!( @@ -1635,7 +1640,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 +1744,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 { @@ -2009,8 +2014,17 @@ pub async fn commit_compaction( return Ok(CompactionMetrics::default()); } - // 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()); + // Address-style results require immediate index remapping unless it is deferred. + let needs_remapping = + !dataset.manifest.uses_stable_row_ids() && !options.defer_index_remap && has_address_style; + + // Confirm there is a remapper before materializing the potentially very large row address map. + let index_remapper = if needs_remapping { + remap_options.create_remapper(dataset).await? + } else { + None + }; // Determine the earliest version at which compaction tasks were planned/executed. // @@ -2034,7 +2048,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() @@ -2084,7 +2097,7 @@ pub async fn commit_compaction( new_fragments: task.new_fragments.clone(), }; - if needs_remapping { + if index_remapper.is_some() { if let Some(row_addrs_bytes) = task.row_addrs { let row_addrs = RoaringTreemap::deserialize_from(&mut Cursor::new(&row_addrs_bytes))?; @@ -2151,8 +2164,7 @@ pub async fn commit_compaction( rewrite_groups.push(rewrite_group); } - let rewritten_indices = if needs_remapping { - let index_remapper = remap_options.create_remapper(dataset)?; + let rewritten_indices = if let Some(index_remapper) = index_remapper { let affected_ids = rewrite_groups .iter() .flat_map(|group| group.old_fragments.iter().map(|frag| frag.id)) @@ -2448,9 +2460,10 @@ mod tests { } } + #[async_trait] impl IndexRemapperOptions for MockIndexRemapper { - fn create_remapper(&self, _: &Dataset) -> Result> { - Ok(Box::new(self.clone())) + async fn create_remapper(&self, _: &Dataset) -> Result>> { + Ok(Some(Box::new(self.clone()))) } } @@ -2979,9 +2992,70 @@ mod tests { } } + #[async_trait] impl IndexRemapperOptions for IgnoreRemap { - fn create_remapper(&self, _: &Dataset) -> Result> { - Ok(Box::new(Self {})) + async fn create_remapper(&self, _: &Dataset) -> Result>> { + Ok(None) + } + } + + #[rstest] + #[case::without_index(false)] + #[case::with_index(true)] + #[tokio::test] + async fn test_row_addrs_only_used_with_remappable_index(#[case] has_index: bool) { + let data = sample_data(); + let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 9_000))], data.schema()); + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 3_000, + data_storage_version: Some(LanceFileVersion::Legacy), + ..Default::default() + }), + ) + .await + .unwrap(); + + if has_index { + create_scalar_index(&mut dataset, "a", false).await; + } + + let options = CompactionOptions { + target_rows_per_fragment: 9_000, + ..Default::default() + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks().len(), 1); + + let mut result = rewrite_files(Cow::Borrowed(&dataset), plan.tasks()[0].clone(), &options) + .await + .unwrap(); + assert_eq!(result.row_addrs.is_some(), has_index); + + if has_index { + let row_addrs_bytes = result + .row_addrs + .as_ref() + .expect("indexed compaction should capture row addresses"); + let row_addrs = + RoaringTreemap::deserialize_from(&mut Cursor::new(row_addrs_bytes)).unwrap(); + assert_eq!(row_addrs.len(), 9_000); + } else { + // Simulate a stale worker result that captured row addresses before the + // dataset no longer needed a remapper. Invalid bytes ensure the commit + // does not attempt to deserialize or materialize the unused map. + result.row_addrs = Some(b"not a roaring treemap".to_vec()); + commit_compaction( + &mut dataset, + vec![result], + Arc::new(DatasetIndexRemapperOptions::default()), + &options, + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 1); } } @@ -3322,17 +3396,10 @@ mod tests { dataset.delete("i < 500").await.unwrap(); dataset2.delete("i < 500").await.unwrap(); - // Create a scalar index to check this is not touched - dataset - .create_index( - &["i"], - IndexType::Scalar, - Some("scalar".into()), - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); + // Create the same scalar index on both datasets so deferred and immediate + // remapping are compared under the same conditions. + create_scalar_index(&mut dataset, "i", false).await; + create_scalar_index(&mut dataset2, "i", false).await; // Verify the initial state - no fragment reuse index should exist let initial_indices = dataset.load_indices().await.unwrap(); diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index aef2cd231fc..682ed696f21 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -61,8 +61,13 @@ pub trait IndexRemapper: Send + Sync { /// /// Currently we don't have any options but we may need options in the future and so we /// want to keep a placeholder +#[async_trait] pub trait IndexRemapperOptions: Send + Sync { - fn create_remapper(&self, dataset: &Dataset) -> Result>; + /// Creates a remapper when the dataset has indices that need row address remapping. + /// + /// Returns `None` when no remappable indices exist, allowing compaction to avoid + /// materializing an unused row address map. + async fn create_remapper(&self, dataset: &Dataset) -> Result>>; } #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] @@ -75,9 +80,10 @@ impl IndexRemapper for IgnoreRemap { } } +#[async_trait] impl IndexRemapperOptions for IgnoreRemap { - fn create_remapper(&self, _: &Dataset) -> Result> { - Ok(Box::new(Self {})) + async fn create_remapper(&self, _: &Dataset) -> Result>> { + Ok(None) } }