diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 34199956c1e..9439c272c72 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -542,25 +542,41 @@ def test_index_remapping_multiple_rewrite_tasks(tmp_path: Path): assert index_frag_ids[0] in frag_ids -def test_defer_index_remap(tmp_path: Path): +@pytest.mark.parametrize("enable_stable_row_ids", [False, True]) +def test_defer_index_remap(tmp_path: Path, enable_stable_row_ids: bool): base_dir = tmp_path / "dataset" data = pa.table({"i": range(6_000), "val": range(6_000)}) - dataset = lance.write_dataset(data, base_dir, max_rows_per_file=1_000) + dataset = lance.write_dataset( + data, + base_dir, + max_rows_per_file=1_000, + enable_stable_row_ids=enable_stable_row_ids, + ) dataset.create_scalar_index("i", "BTREE") options = dict( target_rows_per_fragment=2_000, defer_index_remap=True, num_threads=1 ) dataset.delete("i < 500") + row_ids_before = dataset.to_table(with_row_id=True).column("_rowid").to_pylist() dataset.optimize.compact_files(**options) dataset = lance.dataset(base_dir) indices = dataset.describe_indices() assert any(idx.name == "__lance_frag_reuse" for idx in indices) + row_ids_after = dataset.to_table(with_row_id=True).column("_rowid").to_pylist() + if enable_stable_row_ids: + assert row_ids_after == row_ids_before + # Stable ids 500..1000 collide with the rewritten fragment's addresses; the + # index must still answer with the ids. + assert dataset.to_table(filter="i < 510").num_rows == 10 @pytest.mark.parametrize("use_commit_options", [True, False]) -def test_defer_index_remap_via_commit_options(tmp_path: Path, use_commit_options: bool): +@pytest.mark.parametrize("enable_stable_row_ids", [False, True]) +def test_defer_index_remap_via_commit_options( + tmp_path: Path, use_commit_options: bool, enable_stable_row_ids: bool +): """Compaction.commit respects defer_index_remap passed in options. When options={"defer_index_remap": True} is supplied to Compaction.commit @@ -569,7 +585,12 @@ def test_defer_index_remap_via_commit_options(tmp_path: Path, use_commit_options """ base_dir = tmp_path / f"dataset_commit_opts_{use_commit_options}" data = pa.table({"i": range(6_000), "val": range(6_000)}) - dataset = lance.write_dataset(data, base_dir, max_rows_per_file=1_000) + dataset = lance.write_dataset( + data, + base_dir, + max_rows_per_file=1_000, + enable_stable_row_ids=enable_stable_row_ids, + ) dataset.create_scalar_index("i", "BTREE") dataset.delete("i < 500") diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 7239f3b5991..0de9203cbb7 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -123,6 +123,10 @@ impl ScalarIndex for JsonIndex { .await } + fn results_are_row_addresses(&self) -> bool { + self.target_index.results_are_row_addresses() + } + fn can_remap(&self) -> bool { self.target_index.can_remap() } @@ -1008,14 +1012,13 @@ impl ScalarIndexPlugin for JsonIndexPlugin { index_name: String, index_details: &prost_types::Any, ) -> Option> { - // TODO: Allow return Result here - let registry = self.registry().unwrap(); + let registry = self.registry().ok()?; let json_details = - crate::pb::JsonIndexDetails::decode(index_details.value.as_slice()).unwrap(); - let target_details = json_details.target_details.as_ref().expect_ok().unwrap(); - let target_plugin = registry.get_plugin_by_details(target_details).unwrap(); + crate::pb::JsonIndexDetails::decode(index_details.value.as_slice()).ok()?; + let target_details = json_details.target_details.as_ref()?; + let target_plugin = registry.get_plugin_by_details(target_details).ok()?; // TODO: Use something like ${index_name}_${path} for the index name? Don't have access to path here tho - let target_parser = target_plugin.new_query_parser(index_name, index_details)?; + let target_parser = target_plugin.new_query_parser(index_name, target_details)?; Some(Box::new(JsonQueryParser::new( json_details.path.clone(), target_parser, @@ -1063,6 +1066,35 @@ mod tests { use std::ops::Bound; use std::sync::Arc; + #[test] + fn test_nested_json_query_parser() { + use prost::Message; + let json_type_url = "type.googleapis.com/lance.index.pb.JsonIndexDetails".to_string(); + let btree = prost_types::Any { + type_url: "type.googleapis.com/lance.table.BTreeIndexDetails".to_string(), + value: vec![], + }; + let inner = prost_types::Any { + type_url: json_type_url.clone(), + value: crate::pb::JsonIndexDetails { + path: "a".to_string(), + target_details: Some(btree), + } + .encode_to_vec(), + }; + let outer = prost_types::Any { + type_url: json_type_url, + value: crate::pb::JsonIndexDetails { + path: "b".to_string(), + target_details: Some(inner), + } + .encode_to_vec(), + }; + let registry = crate::registry::IndexPluginRegistry::with_default_plugins(); + let plugin = registry.get_plugin_by_details(&outer).unwrap(); + assert!(plugin.new_query_parser("idx".to_string(), &outer).is_some()); + } + // Note: The old test_detect_json_value_type test has been removed as we now use // JSONB's inherent type information instead of string-based type detection diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 3c67753910d..9680f3f76da 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -56,9 +56,10 @@ pub const FLAG_COVERED_INDEX_METADATA: u64 = 1 << 7; pub const FLAG_MIXED_DATA_FILE_VERSIONS: u64 = 1 << 8; /// The table uses stable row ids and carries a fragment reuse index. /// -/// Reserved ahead of its implementation. This build treats the bit as unknown -/// (see `supported_flags_when`), so a build that knows the flag but not the -/// handling behind it cannot open such a table. +/// Older readers and writers did not expect the combination and could corrupt +/// such a table, so both feature words carry the bit. Set by `build_manifest` +/// whenever both hold, and lifted when either stops holding; see +/// `supported_flags_when`. pub const FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS: u64 = 1 << 9; /// Tagged FRI requires a reader that interprets its mappings and a writer that /// preserves them during maintenance. Legacy-only FRI does not set this bit. @@ -91,12 +92,12 @@ pub fn apply_feature_flags( disable_transaction_file: bool, ) -> Result<()> { // Carried across the reset: a `Manifest` only points at its index section, - // so whether any index declares covering columns is not visible here. `build_manifest` decides it from the index list it is - // committing and sets the bit after calling this; without the carry the - // second call, from `write_manifest_file`, would clear that decision - // immediately before the write. - let covered_index_metadata = (manifest.reader_feature_flags | manifest.writer_feature_flags) - & FLAG_COVERED_INDEX_METADATA; + // so what its indices declare is not visible here. `build_manifest` decides + // these from the index list it is committing and sets the bits after calling + // this; without the carry the second call, from `write_manifest_file`, would + // clear that decision immediately before the write. + let index_derived_flags = (manifest.reader_feature_flags | manifest.writer_feature_flags) + & (FLAG_COVERED_INDEX_METADATA | FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS); let sticky_paired_flags = validated_sticky_paired_flags(manifest)?; // Reset flags @@ -157,8 +158,8 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } - manifest.reader_feature_flags |= covered_index_metadata; - manifest.writer_feature_flags |= covered_index_metadata; + manifest.reader_feature_flags |= index_derived_flags; + manifest.writer_feature_flags |= index_derived_flags; manifest.reader_feature_flags |= sticky_paired_flags; manifest.writer_feature_flags |= sticky_paired_flags; @@ -198,17 +199,22 @@ fn mark_supported(flags: &mut u64, flag: u64, feature_enabled: bool) { } /// The feature-flag bits this build understands, given whether overlay support -/// is enabled. Split out from [`supported_flags`] so the policy is testable -/// without toggling the build profile or environment. -fn supported_flags_when(overlay_enabled: bool) -> u64 { +/// is enabled and whether a fragment reuse index on a stable-row-id table is +/// supported. Split out from [`supported_flags`] so the policy is testable +/// without toggling the build profile or environment, and so the compatibility +/// test can show what a build without that support does with such a table. +fn supported_flags_when(overlay_enabled: bool, frag_reuse_with_stable_row_ids: bool) -> u64 { let mut supported = FLAG_UNKNOWN - 1; mark_supported( &mut supported, FLAG_UNSTABLE_DATA_OVERLAY_FILES, overlay_enabled, ); - // Reserved, not implemented: see the flag's doc comment. - mark_supported(&mut supported, FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, false); + mark_supported( + &mut supported, + FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + frag_reuse_with_stable_row_ids, + ); // Bit 10 now falls below the unknown boundary, so keep tagged FRI refused // until its reader/writer handling lands. mark_supported(&mut supported, FLAG_FRAGMENT_REUSE_INDEX, false); @@ -216,7 +222,7 @@ fn supported_flags_when(overlay_enabled: bool) -> u64 { } fn supported_flags() -> u64 { - supported_flags_when(data_overlay_files_enabled()) + supported_flags_when(data_overlay_files_enabled(), true) } pub fn can_read_dataset(reader_flags: u64) -> bool { @@ -314,18 +320,25 @@ mod tests { use super::*; use crate::format::BasePath; - /// Reserved ahead of its implementation: refused for reading and writing - /// until the handling lands, so a build from the gap cannot open the table. + /// A build without support for the combination must refuse the table for + /// reading and writing, or it could corrupt it; this build accepts it. #[test] - fn test_frag_reuse_with_stable_row_ids_flag_is_reserved_not_supported() { + fn test_frag_reuse_with_stable_row_ids_flag_gating() { use crate::format::{DataStorageFormat, Manifest}; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use lance_core::datatypes::Schema; use std::collections::HashMap; use std::sync::Arc; - assert!(!can_read_dataset(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS)); - assert!(!can_write_dataset(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS)); + let without_support = supported_flags_when(true, false); + assert_eq!( + without_support & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0, + "a build without the implementation must not accept the flag" + ); + assert_ne!(without_support & FLAG_STABLE_ROW_IDS, 0); + assert!(can_read_dataset(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS)); + assert!(can_write_dataset(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS)); let arrow_schema = ArrowSchema::new(vec![ArrowField::new( "id", @@ -340,14 +353,13 @@ mod tests { ); manifest.reader_feature_flags = FLAG_STABLE_ROW_IDS | FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS; manifest.writer_feature_flags = FLAG_STABLE_ROW_IDS | FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS; - assert!(matches!( - ensure_can_read_manifest(&manifest).unwrap_err(), - Error::NotSupported { .. } - )); - assert!(matches!( - ensure_can_write_manifest(&manifest).unwrap_err(), - Error::NotSupported { .. } - )); + ensure_can_read_manifest(&manifest).unwrap(); + ensure_can_write_manifest(&manifest).unwrap(); + assert_ne!( + manifest.reader_feature_flags & !without_support, + 0, + "the same manifest is refused without the implementation" + ); } #[test] @@ -378,12 +390,12 @@ mod tests { fn test_data_overlay_flag_release_gating() { // Release default (overlays disabled): the overlay flag is treated as // unknown so the dataset is refused, while other known flags still pass. - let supported = supported_flags_when(false); + let supported = supported_flags_when(false, true); assert_eq!(supported & FLAG_UNSTABLE_DATA_OVERLAY_FILES, 0); assert_eq!(FLAG_DELETION_FILES & !supported, 0); assert_ne!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); // Enabled (debug or env opt-in): the overlay flag is understood. - let supported = supported_flags_when(true); + let supported = supported_flags_when(true, true); assert_eq!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); } diff --git a/rust/lance-table/src/format.rs b/rust/lance-table/src/format.rs index 1c9e0e37c8c..9465a856cf6 100644 --- a/rust/lance-table/src/format.rs +++ b/rust/lance-table/src/format.rs @@ -16,7 +16,10 @@ pub use crate::rowids::version::{ RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, }; pub use fragment::*; -pub use index::{IndexFile, IndexMetadata, index_metadata_codec, list_index_files_with_sizes}; +pub use index::{ + IdentifierDomain, IndexFile, IndexMetadata, MAX_JSON_INDEX_NESTING, index_metadata_codec, + list_index_files_with_sizes, +}; pub use manifest::{ BasePath, DETACHED_VERSION_MASK, DataStorageFormat, Manifest, ManifestBuildConfig, diff --git a/rust/lance-table/src/format/index.rs b/rust/lance-table/src/format/index.rs index bfd195a2a7f..3f0ccecc88e 100644 --- a/rust/lance-table/src/format/index.rs +++ b/rust/lance-table/src/format/index.rs @@ -17,6 +17,7 @@ use uuid::Uuid; use super::pb; use lance_core::cache::{CacheEntryReader, CacheEntryWriter}; use lance_core::{Error, Result}; +use prost::Message; /// Metadata about a single file within an index segment. #[derive(Debug, Clone, PartialEq, DeepSizeOf)] @@ -93,6 +94,77 @@ pub struct IndexMetadata { pub files: Option>, } +/// Full type name after the last `/`, case-insensitive, as the index registry matches. +fn type_name_is(type_url: &str, full_name: &str) -> bool { + type_url + .rsplit_once('/') + .map_or(type_url, |(_, name)| name) + .eq_ignore_ascii_case(full_name) +} + +/// Deepest chain of JSON index wrappers this build follows. +pub const MAX_JSON_INDEX_NESTING: usize = 4; + +/// The identifiers an index stores, which decides whether a fragment reuse index may +/// remap them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IdentifierDomain { + RowAddress, + StableRowId, + /// A detail type this build does not recognize; treat conservatively. + Unknown, +} + +const ROW_ADDR_DETAILS: &[&str] = &[ + "lance.table.ZoneMapIndexDetails", + // Lance 0.36 released the zone map details in the index package. + "lance.index.pb.ZoneMapIndexDetails", + "lance.index.pb.BloomFilterIndexDetails", + "lance.index.pb.FMIndexDetails", +]; + +/// Index types whose entries are the `_rowid` column. +const ROW_ID_DETAILS: &[&str] = &[ + "lance.table.BTreeIndexDetails", + "lance.table.BitmapIndexDetails", + "lance.table.LabelListIndexDetails", + "lance.table.NGramIndexDetails", + "lance.table.InvertedIndexDetails", + // Lance 0.36 released these in the index package. + "lance.index.pb.BTreeIndexDetails", + "lance.index.pb.BitmapIndexDetails", + "lance.index.pb.LabelListIndexDetails", + "lance.index.pb.NGramIndexDetails", + "lance.index.pb.InvertedIndexDetails", + "lance.index.pb.RTreeIndexDetails", + "lance.index.pb.VectorIndexDetails", + "lance.index.VectorIndexDetails", +]; + +fn type_url_domain(type_url: &str) -> IdentifierDomain { + if ROW_ADDR_DETAILS + .iter() + .any(|full_name| type_name_is(type_url, full_name)) + { + IdentifierDomain::RowAddress + } else if ROW_ID_DETAILS + .iter() + .any(|full_name| type_name_is(type_url, full_name)) + { + IdentifierDomain::StableRowId + } else { + IdentifierDomain::Unknown + } +} + +/// `JsonIndexDetails.target_details` (field 2 in `index.proto`), decoded here because +/// `lance-index` depends on this crate. +#[derive(Clone, PartialEq, prost::Message)] +struct JsonIndexDetailsTarget { + #[prost(message, optional, tag = "2")] + target_details: Option, +} + impl IndexMetadata { pub fn effective_fragment_bitmap( &self, @@ -135,24 +207,62 @@ impl IndexMetadata { } /// True when the index reports matches as physical row addresses rather than row ids - /// (`ScalarIndex::results_are_row_addresses`). - /// - /// Such an index cannot follow its data through a rewrite: the addresses it stores - /// name fragments and offsets, and neither kind supports remap. + /// (`ScalarIndex::results_are_row_addresses`), judged from the outer details type only. + #[deprecated( + since = "13.0.0", + note = "ignores the target of a JSON index and cannot report an unknown type; use identifier_domain" + )] pub fn results_are_row_addrs(&self) -> bool { self.index_details.as_ref().is_some_and(|details| { - let is_fm = details - .type_url - .rsplit_once('/') - .is_some_and(|(_, details_type_name)| { - details_type_name.eq_ignore_ascii_case("lance.index.pb.FMIndexDetails") - }); - details.type_url.ends_with("ZoneMapIndexDetails") - || details.type_url.ends_with("BloomFilterIndexDetails") - || is_fm + type_url_domain(&details.type_url) == IdentifierDomain::RowAddress }) } + /// The domain the index type declares for a stable-row-id table. Metadata without + /// details is a legacy row-id index. A JSON index answers in its target's domain, so + /// JSON details that do not decode or name no target are an error rather than a guess. + pub fn declared_identifier_domain(&self) -> Result { + let Some(outer) = self.index_details.as_ref() else { + return Ok(IdentifierDomain::StableRowId); + }; + let mut details = outer.as_ref().clone(); + for _ in 0..=MAX_JSON_INDEX_NESTING { + if !type_name_is(&details.type_url, "lance.index.pb.JsonIndexDetails") { + return Ok(type_url_domain(&details.type_url)); + } + let json_details = + JsonIndexDetailsTarget::decode(details.value.as_slice()).map_err(|err| { + Error::corrupt_file_named( + "index metadata", + format!( + "index {} ({}) has JsonIndexDetails that do not decode: {err}", + self.name, self.uuid + ), + ) + })?; + details = json_details.target_details.ok_or_else(|| { + Error::corrupt_file_named( + "index metadata", + format!( + "index {} ({}) has JsonIndexDetails without target_details", + self.name, self.uuid + ), + ) + })?; + } + Ok(IdentifierDomain::Unknown) + } + + /// The identifiers this index stores. Without stable row ids every index stores row + /// addresses; with them the type decides, and remapping stable row ids as if they + /// were addresses rewrites unrelated rows. + pub fn identifier_domain(&self, uses_stable_row_ids: bool) -> Result { + if !uses_stable_row_ids { + return Ok(IdentifierDomain::RowAddress); + } + self.declared_identifier_domain() + } + /// The prefix of [`Self::fields`] this index is keyed on, with the carried /// columns of [`Self::covering_fields`] removed. /// @@ -642,21 +752,166 @@ mod tests { } } + fn json_details_over(target_type_url: &str) -> Vec { + json_details_over_any(prost_types::Any { + type_url: target_type_url.to_string(), + value: Vec::new(), + }) + } + + fn json_details_over_any(target: prost_types::Any) -> Vec { + JsonIndexDetailsTarget { + target_details: Some(target), + } + .encode_to_vec() + } + + fn nested_json(depth: usize, innermost_type_url: &str) -> prost_types::Any { + let mut details = prost_types::Any { + type_url: innermost_type_url.to_string(), + value: Vec::new(), + }; + for _ in 0..depth { + details = prost_types::Any { + type_url: "type.googleapis.com/lance.index.pb.JsonIndexDetails".to_string(), + value: json_details_over_any(details), + }; + } + details + } + #[rstest] - #[case::zone_map("type.googleapis.com/lance.table.ZoneMapIndexDetails", true)] - #[case::bloom_filter("type.googleapis.com/lance.index.pb.BloomFilterIndexDetails", true)] - #[case::fm("type.googleapis.com/lance.index.pb.FMIndexDetails", true)] - #[case::fm_case_insensitive("type.googleapis.com/LANCE.INDEX.PB.FMINDEXDETAILS", true)] - #[case::foreign_fm_terminal_name("type.googleapis.com/example.FMIndexDetails", false)] - #[case::btree("type.googleapis.com/lance.table.BTreeIndexDetails", false)] - fn test_results_are_row_addrs(#[case] type_url: &str, #[case] expected: bool) { + #[case::zone_map("type.googleapis.com/lance.table.ZoneMapIndexDetails", vec![], IdentifierDomain::RowAddress)] + #[case::zone_map_v036_package("type.googleapis.com/lance.index.pb.ZoneMapIndexDetails", vec![], IdentifierDomain::RowAddress)] + #[case::bloom_filter("type.googleapis.com/lance.index.pb.BloomFilterIndexDetails", vec![], IdentifierDomain::RowAddress)] + #[case::fm("type.googleapis.com/lance.index.pb.FMIndexDetails", vec![], IdentifierDomain::RowAddress)] + #[case::fm_case_insensitive("type.googleapis.com/LANCE.INDEX.PB.FMINDEXDETAILS", vec![], IdentifierDomain::RowAddress)] + #[case::btree("type.googleapis.com/lance.table.BTreeIndexDetails", vec![], IdentifierDomain::StableRowId)] + #[case::btree_v036_package("type.googleapis.com/lance.index.pb.BTreeIndexDetails", vec![], IdentifierDomain::StableRowId)] + #[case::rtree("type.googleapis.com/lance.index.pb.RTreeIndexDetails", vec![], IdentifierDomain::StableRowId)] + #[case::vector("type.googleapis.com/lance.index.pb.VectorIndexDetails", vec![], IdentifierDomain::StableRowId)] + #[case::vector_legacy_package("type.googleapis.com/lance.index.VectorIndexDetails", vec![], IdentifierDomain::StableRowId)] + #[case::foreign_fm_terminal_name("type.googleapis.com/example.FMIndexDetails", vec![], IdentifierDomain::Unknown)] + #[case::foreign_zone_map_terminal_name("type.googleapis.com/example.ZoneMapIndexDetails", vec![], IdentifierDomain::Unknown)] + #[case::future_type("type.googleapis.com/lance.table.FutureAddressIndexDetails", vec![], IdentifierDomain::Unknown)] + #[case::json_over_zone_map( + "type.googleapis.com/lance.index.pb.JsonIndexDetails", + json_details_over("type.googleapis.com/lance.table.ZoneMapIndexDetails"), + IdentifierDomain::RowAddress + )] + #[case::json_over_v036_zone_map( + "type.googleapis.com/lance.index.pb.JsonIndexDetails", + json_details_over("type.googleapis.com/lance.index.pb.ZoneMapIndexDetails"), + IdentifierDomain::RowAddress + )] + #[case::json_upper_case_over_zone_map( + "type.googleapis.com/LANCE.INDEX.PB.JSONINDEXDETAILS", + json_details_over("type.googleapis.com/LANCE.TABLE.ZONEMAPINDEXDETAILS"), + IdentifierDomain::RowAddress + )] + #[case::json_over_btree( + "type.googleapis.com/lance.index.pb.JsonIndexDetails", + json_details_over("type.googleapis.com/lance.table.BTreeIndexDetails"), + IdentifierDomain::StableRowId + )] + #[case::json_over_unknown_target( + "type.googleapis.com/lance.index.pb.JsonIndexDetails", + json_details_over("type.googleapis.com/lance.table.FutureAddressIndexDetails"), + IdentifierDomain::Unknown + )] + fn test_identifier_domain( + #[case] type_url: &str, + #[case] value: Vec, + #[case] expected: IdentifierDomain, + ) { let mut metadata = index_metadata_with(vec![0], vec![]); metadata.index_details = Some(Arc::new(prost_types::Any { type_url: type_url.to_string(), + value, + })); + + assert_eq!(metadata.declared_identifier_domain().unwrap(), expected); + assert_eq!(metadata.identifier_domain(true).unwrap(), expected); + assert_eq!( + metadata.identifier_domain(false).unwrap(), + IdentifierDomain::RowAddress + ); + } + + #[test] + fn test_identifier_domain_follows_nested_json_up_to_the_cap() { + let zone_map = "type.googleapis.com/lance.table.ZoneMapIndexDetails"; + let mut metadata = index_metadata_with(vec![0], vec![]); + metadata.index_details = Some(Arc::new(nested_json(MAX_JSON_INDEX_NESTING, zone_map))); + assert_eq!( + metadata.declared_identifier_domain().unwrap(), + IdentifierDomain::RowAddress + ); + metadata.index_details = Some(Arc::new(nested_json(MAX_JSON_INDEX_NESTING + 1, zone_map))); + assert_eq!( + metadata.declared_identifier_domain().unwrap(), + IdentifierDomain::Unknown + ); + } + + #[test] + fn test_identifier_domain_without_details_is_legacy_row_id() { + let metadata = index_metadata_with(vec![0], vec![]); + assert_eq!( + metadata.identifier_domain(true).unwrap(), + IdentifierDomain::StableRowId + ); + assert_eq!( + metadata.identifier_domain(false).unwrap(), + IdentifierDomain::RowAddress + ); + } + + /// The deprecated method keeps judging by the outer type only. + #[test] + #[allow(deprecated)] + fn test_results_are_row_addrs_ignores_json_target() { + let mut metadata = index_metadata_with(vec![0], vec![]); + metadata.index_details = Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.pb.JsonIndexDetails".to_string(), + value: json_details_over("type.googleapis.com/lance.table.ZoneMapIndexDetails"), + })); + assert!(!metadata.results_are_row_addrs()); + assert_eq!( + metadata.declared_identifier_domain().unwrap(), + IdentifierDomain::RowAddress + ); + + metadata.index_details = Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.ZoneMapIndexDetails".to_string(), value: Vec::new(), })); + assert!(metadata.results_are_row_addrs()); + } - assert_eq!(metadata.results_are_row_addrs(), expected); + /// The domain decides whether stored identifiers get rewritten, so JSON details + /// that establish no target are refused instead of read as "stable row ids". + #[rstest] + #[case::undecodable(vec![0xff], "do not decode")] + #[case::no_target(vec![], "without target_details")] + fn test_results_are_row_addrs_rejects_incomplete_json_details( + #[case] value: Vec, + #[case] expected_message: &str, + ) { + let mut metadata = index_metadata_with(vec![0], vec![]); + metadata.index_details = Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.pb.JsonIndexDetails".to_string(), + value, + })); + + let error = metadata.declared_identifier_domain().unwrap_err(); + assert!(matches!(error, Error::CorruptFile { .. }), "{error}"); + assert!(error.to_string().contains(expected_message), "{error}"); + assert!(metadata.identifier_domain(true).is_err()); + assert_eq!( + metadata.identifier_domain(false).unwrap(), + IdentifierDomain::RowAddress + ); } #[rstest] diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 6ea623c60da..dfc0c2b09d7 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -19,7 +19,9 @@ use std::ops::Range; use std::sync::Arc; use super::{Fragment, InlineRowIds, RowIdMeta}; -use crate::feature_flags::{FLAG_COVERED_INDEX_METADATA, STICKY_PAIRED_FLAGS}; +use crate::feature_flags::{ + FLAG_COVERED_INDEX_METADATA, FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, STICKY_PAIRED_FLAGS, +}; use crate::feature_flags::{FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag}; use crate::format::fragment::DataFileFieldInterner; use crate::format::pb; @@ -283,13 +285,20 @@ impl Manifest { // wholesale -- `covering_fields` included -- a build that predates // covering could then open it and read carried columns as keyed ones. // Kept unconditionally rather than derived from the cloned indexes: - // over-fencing a clone is harmless, under-fencing one is not. + // over-fencing a clone is harmless, under-fencing one is not. The + // same holds for the stable-row-id fragment reuse fence: the clone + // copies both the stable row ids and the fragment reuse index, so it + // needs the same protection from older readers and writers. // Sticky capabilities are also retained because the clone keeps the // source file identities that require them. reader_feature_flags: self.reader_feature_flags - & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS), + & (FLAG_COVERED_INDEX_METADATA + | FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS + | STICKY_PAIRED_FLAGS), writer_feature_flags: self.writer_feature_flags - & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS), + & (FLAG_COVERED_INDEX_METADATA + | FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS + | STICKY_PAIRED_FLAGS), max_fragment_id: self.max_fragment_id, transaction_file: Some(transaction_file), transaction_section: None, diff --git a/rust/lance-table/src/system_index/frag_reuse.rs b/rust/lance-table/src/system_index/frag_reuse.rs index 56a09722edf..4ab11192600 100644 --- a/rust/lance-table/src/system_index/frag_reuse.rs +++ b/rust/lance-table/src/system_index/frag_reuse.rs @@ -647,6 +647,33 @@ mod tests { ); } + /// Stable row ids start at 0, so fragment 0's addresses are numerically the + /// first stable ids. The remapper cannot tell them apart: applied to + /// stable-id entries it rewrites them, which is why index loading gates on + /// the identifier domain. + #[test] + fn test_compact_fri_rewrites_colliding_stable_row_ids() { + let details = FragReuseIndexDetails { + versions: vec![FragReuseVersion { + dataset_version: 1, + groups: vec![FragReuseGroup { + changed_row_addrs: serialize_changed([addr(0, 0), addr(0, 2)]), + old_frags: vec![digest(0, 3)], + new_frags: vec![digest(7, 2)], + }], + }], + }; + let fri = CompactFragReuseIndex::try_new(Uuid::new_v4(), details).unwrap(); + + // Stable ids 0, 1, 2 happen to equal fragment 0's addresses. + let mut stable_ids = vec![Some(0), Some(1), Some(2), Some(3)]; + fri.remap_row_ids_in_place(&mut stable_ids); + assert_eq!( + stable_ids, + vec![Some(addr(7, 0)), None, Some(addr(7, 1)), Some(3)] + ); + } + #[test] fn test_compact_fri_rejects_invalid_changed_row_bitmap() { let details = FragReuseIndexDetails { diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index 06cf9648a88..c59ee8385f6 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -10,7 +10,7 @@ //! returns stale rows from the index -- so each rule here is paired with a test. use crate::format::overlay::staleness::collect_overlay_stale_frags; -use crate::format::{Fragment, IndexMetadata}; +use crate::format::{Fragment, IdentifierDomain, IndexMetadata}; use crate::system_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use crate::system_index::is_system_index; use crate::transaction::{RewriteGroup, RewrittenIndex, Transaction}; @@ -27,6 +27,7 @@ impl Transaction { fields_for_preserving_frag_bitmap: &[u32], original_overlaid_frags: &HashMap, schema: &Schema, + index_domain: &dyn Fn(&IndexMetadata) -> Result, ) -> Result<()> { if pure_update_frag_ids.is_empty() { return Ok(()); @@ -39,7 +40,7 @@ impl Transaction { for index in indices.iter_mut() { // Physical row addresses cannot follow moved rows into a new fragment. // Leave that fragment uncovered so the scanner reads it directly. - if index.results_are_row_addrs() { + if index_domain(index)? != IdentifierDomain::StableRowId { continue; } let index_covers_modified_field = index.fields.iter().any(|field_id| { diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 5551bd36e64..701af9ac03f 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -11,13 +11,14 @@ //! metadata it stamps, the validation that runs before it. use crate::feature_flags::{ - FLAG_COVERED_INDEX_METADATA, FLAG_STABLE_ROW_IDS, apply_feature_flags, - ensure_can_read_manifest, ensure_can_write_manifest, inherit_sticky_feature_flags, + FLAG_COVERED_INDEX_METADATA, FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, FLAG_STABLE_ROW_IDS, + apply_feature_flags, ensure_can_read_manifest, ensure_can_write_manifest, + inherit_sticky_feature_flags, }; use crate::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{ - DataFile, DataStorageFormat, Fragment, IndexMetadata, Manifest, ManifestBuildConfig, RowIdMeta, - overlay::DataOverlayFile, + DataFile, DataStorageFormat, Fragment, IdentifierDomain, IndexMetadata, Manifest, + ManifestBuildConfig, RowIdMeta, overlay::DataOverlayFile, }; use crate::io::{ commit::CommitHandler, @@ -119,6 +120,15 @@ fn apply_range_segments( Ok(()) } +/// Without version knowledge, any index with details may store addresses. +fn conservative_index_domain(index: &IndexMetadata) -> Result { + Ok(if index.index_details.is_some() { + IdentifierDomain::Unknown + } else { + IdentifierDomain::StableRowId + }) +} + impl Transaction { pub(super) fn fragments_with_ids<'a, T>( new_fragments: T, @@ -450,6 +460,9 @@ impl Transaction { /// Create a new manifest from the current manifest and the transaction. /// /// `current_manifest` should only be None if the dataset does not yet exist. + /// Classifies existing indexes without version knowledge, so a rewrite on a + /// stable-row-id table may withdraw their coverage; see + /// [`Self::build_manifest_with_index_domain`]. pub fn build_manifest( &self, current_manifest: Option<&Manifest>, @@ -479,6 +492,29 @@ impl Transaction { transaction_file_path: &str, config: &ManifestBuildConfig, read_version_state: Option>, + ) -> Result<(Manifest, Vec)> { + self.build_manifest_with_index_domain( + current_manifest, + current_indices, + transaction_file_path, + config, + read_version_state, + &conservative_index_domain, + ) + } + + /// [`Self::build_manifest_with_read_version`] with the caller's classification of + /// each existing index on a stable-row-id table. Only [`IdentifierDomain::StableRowId`] + /// coverage follows a rewrite; the caller is expected to return + /// [`IdentifierDomain::Unknown`] for any index version it cannot read. + pub fn build_manifest_with_index_domain( + &self, + current_manifest: Option<&Manifest>, + current_indices: Vec, + transaction_file_path: &str, + config: &ManifestBuildConfig, + read_version_state: Option>, + index_domain: &dyn Fn(&IndexMetadata) -> Result, ) -> Result<(Manifest, Vec)> { if config.use_stable_row_ids && config.migration_next_row_id.is_none() @@ -844,6 +880,7 @@ impl Transaction { fields_for_preserving_frag_bitmap, &original_overlaid_frags, &schema, + index_domain, )?; } @@ -906,23 +943,25 @@ impl Transaction { )?; if next_row_id.is_some() { - // We can re-use indices, but need to rewrite the fragment bitmaps + // We can re-use indices, but need to rewrite the fragment bitmaps. + // This holds with a fragment reuse index too: its load-time + // coverage remap finds no rewritten fragment left in these + // bitmaps and leaves them alone. debug_assert!(rewritten_indices.is_empty()); for index in final_indices.iter_mut() { - let results_are_row_addrs = index.results_are_row_addrs(); + // Not data coverage: the fragment reuse index's bitmap is + // its chain's output, and the eager path leaves it alone too. + if index.name == FRAG_REUSE_INDEX_NAME { + continue; + } + // Only stable-row-id entries survive a rewrite; anything else + // loses the rewritten fragments and the scanner reads them. + let follows_rewrite = index_domain(index)? == IdentifierDomain::StableRowId; if let Some(fragment_bitmap) = &mut index.fragment_bitmap { - *fragment_bitmap = if results_are_row_addrs { - // Stable row ids survive a rewrite, so a row-id-domain index - // can simply follow its data to the new fragments. An - // address-domain index cannot: its stored addresses point into - // the fragments the rewrite dropped. Claiming coverage of the - // new fragments would make it answer queries with addresses - // that no longer resolve, so drop the rewritten fragments from - // its coverage instead and let the scanner fall back to a full - // scan for them. - Self::drop_rewritten_fragments(fragment_bitmap, groups) - } else { + *fragment_bitmap = if follows_rewrite { Self::recalculate_fragment_bitmap(fragment_bitmap, groups)? + } else { + Self::drop_rewritten_fragments(fragment_bitmap, groups) }; } } @@ -1457,6 +1496,17 @@ impl Transaction { manifest.reader_feature_flags |= FLAG_COVERED_INDEX_METADATA; manifest.writer_feature_flags |= FLAG_COVERED_INDEX_METADATA; } + // Derived the same way. Older readers and writers did not expect stable + // row ids and a fragment reuse index together and could corrupt such a + // table, so both words are fenced. + if manifest.uses_stable_row_ids() + && final_indices + .iter() + .any(|index| index.name == FRAG_REUSE_INDEX_NAME) + { + manifest.reader_feature_flags |= FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS; + manifest.writer_feature_flags |= FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS; + } if let Some(current_manifest) = current_manifest { inherit_sticky_feature_flags(&mut manifest, current_manifest)?; @@ -2428,6 +2478,168 @@ mod tests { assert_eq!(seq.version_at(4).unwrap(), 2); } + /// A deferred rewrite on a stable-row-id table: coverage moves by identifier + /// domain, the fragment reuse index entry is replaced rather than recomputed + /// (its bitmap is the chain's output and may straddle the group), and the + /// compatibility flag follows the index's presence. + #[test] + fn rewrite_build_manifest_with_frag_reuse_index_and_stable_row_ids() { + use crate::feature_flags::{FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, FLAG_STABLE_ROW_IDS}; + use crate::system_index::frag_reuse::FRAG_REUSE_INDEX_NAME; + use uuid::Uuid; + + let arrow_schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let lance_schema = LanceSchema::try_from(&arrow_schema).unwrap(); + let fragment = |id: u64, row_ids: &[u64]| Fragment { + id, + files: vec![DataFile::new( + format!("{id}.lance"), + vec![0], + vec![0], + LanceFileVersion::Stable.resolve(), + None, + None, + )], + overlays: vec![], + deletion_file: None, + row_id_meta: Some(RowIdMeta::Inline( + write_row_ids(&RowIdSequence::from(row_ids)).into(), + )), + physical_rows: Some(row_ids.len()), + last_updated_at_version_meta: None, + created_at_version_meta: None, + }; + let index = |name: &str, type_name: &str, coverage: &[u32]| IndexMetadata { + uuid: Uuid::new_v4(), + name: name.to_string(), + fields: vec![0], + covering_fields: vec![], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter(coverage.iter().copied())), + index_details: Some(Arc::new(prost_types::Any { + type_url: format!("type.googleapis.com/{type_name}"), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + }; + let mut manifest = Manifest::new( + lance_schema, + Arc::new(vec![ + fragment(0, &[100, 101]), + fragment(1, &[102, 103]), + fragment(2, &[104]), + ]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + manifest.reader_feature_flags |= FLAG_STABLE_ROW_IDS; + manifest.writer_feature_flags |= FLAG_STABLE_ROW_IDS; + manifest.next_row_id = 105; + + let indices = vec![ + index("stable", "lance.table.BTreeIndexDetails", &[0, 1, 2]), + index("addrs", "lance.table.ZoneMapIndexDetails", &[0, 1, 2]), + // Output of an earlier deferred rewrite that produced fragment 1 alone, + // so it straddles the group below. + index( + FRAG_REUSE_INDEX_NAME, + "lance.table.FragmentReuseIndexDetails", + &[1], + ), + ]; + let new_frag_reuse_index = index( + FRAG_REUSE_INDEX_NAME, + "lance.table.FragmentReuseIndexDetails", + &[3], + ); + let rewrite = |frag_reuse_index: Option| { + Transaction::new( + manifest.version, + Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![fragment(0, &[100, 101]), fragment(1, &[102, 103])], + new_fragments: vec![fragment(3, &[100, 101, 102, 103])], + }], + rewritten_indices: vec![], + frag_reuse_index, + }, + None, + ) + }; + let mut config = default_build_config(); + config.use_stable_row_ids = true; + + let (out, out_indices) = rewrite(Some(new_frag_reuse_index.clone())) + .build_manifest_with_index_domain( + Some(&manifest), + indices.clone(), + "txn", + &config, + None, + &|index: &IndexMetadata| index.identifier_domain(true), + ) + .unwrap(); + let coverage = |name: &str| { + out_indices + .iter() + .find(|index| index.name == name) + .unwrap() + .fragment_bitmap + .clone() + .unwrap() + }; + assert_eq!(coverage("stable"), RoaringBitmap::from_iter([2, 3])); + assert_eq!(coverage("addrs"), RoaringBitmap::from_iter([2])); + assert_eq!( + out_indices + .iter() + .find(|index| index.name == FRAG_REUSE_INDEX_NAME) + .unwrap() + .uuid, + new_frag_reuse_index.uuid + ); + assert_ne!( + out.reader_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + assert_ne!( + out.writer_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + + // Without version knowledge, the plain entry point withdraws even a known + // type's coverage. + let (_, fallback_indices) = rewrite(Some(new_frag_reuse_index)) + .build_manifest(Some(&manifest), indices.clone(), "txn", &config) + .unwrap(); + assert_eq!( + fallback_indices + .iter() + .find(|index| index.name == "stable") + .unwrap() + .fragment_bitmap + .clone() + .unwrap(), + RoaringBitmap::from_iter([2]) + ); + + // No fragment reuse index in the result: the flag is not set. + let (out, _) = rewrite(None) + .build_manifest(Some(&manifest), indices[..2].to_vec(), "txn", &config) + .unwrap(); + assert_eq!( + out.reader_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + assert_eq!( + out.writer_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + } + #[test] fn merge_build_manifest_skips_refresh_when_carry_forward_stable_row_ids() { use crate::feature_flags::FLAG_STABLE_ROW_IDS; diff --git a/rust/lance/src/dataset/index/frag_reuse.rs b/rust/lance/src/dataset/index/frag_reuse.rs index ce69419b6ea..2a234b56486 100644 --- a/rust/lance/src/dataset/index/frag_reuse.rs +++ b/rust/lance/src/dataset/index/frag_reuse.rs @@ -5,15 +5,15 @@ use std::sync::Arc; use crate::Dataset; use crate::dataset::transaction::{Operation, Transaction}; -use crate::index::DatasetIndexInternalExt; use crate::index::frag_reuse::{build_frag_reuse_index_metadata, load_frag_reuse_index_details}; +use crate::index::{DatasetIndexInternalExt, effective_identifier_domain}; use lance_core::{Error, Result}; use lance_index::frag_reuse::{ CompactFragReuseIndex, FRAG_REUSE_INDEX_NAME, FragReuseIndexDetails, FragReuseVersion, }; use lance_index::is_system_index; use lance_index::metrics::NoOpMetricsCollector; -use lance_table::format::IndexMetadata; +use lance_table::format::{IdentifierDomain, IndexMetadata}; use lance_table::io::manifest::read_manifest_indexes; use log::warn; use roaring::RoaringBitmap; @@ -43,8 +43,10 @@ impl Dataset { /// * Mapped destinations are not checked against the manifest and can be /// stale, for example after every row of the destination fragment is deleted. /// * Not every compaction records an FRI: it requires `defer_index_remap`, - /// fresh index-free tables do not receive one automatically, and datasets - /// with stable row ids reject the option. + /// and fresh index-free tables do not receive one automatically. + /// * With stable row ids the FRI still describes physical addresses. A table + /// version with both sets reader and writer feature flags that older Lance + /// versions, which could corrupt such a table, refuse. /// * Says nothing about deletion files, source-value changes, or whether an /// address belongs to this table or branch. /// @@ -84,6 +86,11 @@ impl Dataset { /// 2. it is at or past the reuse version's dataset version and no old fragment /// in the version is still in its bitmap. A missing bitmap counts as caught /// up, else the version could never be cleaned up. +/// 3. it stores stable row ids, as indices on stable-row-id tables do today, +/// which a rewrite does not move, and no old fragment in the version is still +/// in its bitmap. Its data never needs the mapping; only coverage committed +/// against fragments the version rewrote (an index created concurrently with +/// the compaction) still does. /// /// Note that there could be a race condition that an index is being added during the cleanup, /// This will make that specific index not efficient until the next reindex, @@ -124,12 +131,23 @@ pub async fn cleanup_frag_reuse_index(dataset: &mut Dataset) -> lance_core::Resu let chain_frag_bitmap = reuse_chain_frag_bitmap(&frag_reuse_details.versions); + let uses_stable_row_ids = dataset.manifest.uses_stable_row_ids(); + let mut index_domains = Vec::with_capacity(indices.len()); + for idx in indices.iter() { + index_domains.push(( + idx, + effective_identifier_domain(idx, uses_stable_row_ids)? == IdentifierDomain::StableRowId, + )); + } + let mut retained_versions = Vec::new(); let mut fragment_bitmaps = RoaringBitmap::new(); for version in frag_reuse_details.versions.iter() { - let check_results = indices + let check_results = index_domains .iter() - .map(|idx| is_index_remap_caught_up(version, idx, &chain_frag_bitmap)) + .map(|(idx, is_data_domain_stable)| { + is_index_remap_caught_up(version, idx, &chain_frag_bitmap, *is_data_domain_stable) + }) .collect::>(); if check_results @@ -195,6 +213,7 @@ fn is_index_remap_caught_up( frag_reuse_version: &FragReuseVersion, index_meta: &IndexMetadata, chain_frag_bitmap: &RoaringBitmap, + is_data_domain_stable: bool, ) -> lance_core::Result { if is_system_index(index_meta) { return Ok(true); @@ -211,7 +230,7 @@ fn is_index_remap_caught_up( return Ok(true); } - if index_meta.dataset_version < frag_reuse_version.dataset_version { + if !is_data_domain_stable && index_meta.dataset_version <= frag_reuse_version.dataset_version { return Ok(false); } @@ -314,23 +333,26 @@ mod tests { // Non-covering, stale version: touches none of the rewritten frags, so // caught up despite version 5 < 10 (the case the old gate got wrong). assert_true!( - is_index_remap_caught_up(&version, &index_covering(5, &[1, 2, 3]), &chain).unwrap() + is_index_remap_caught_up(&version, &index_covering(5, &[1, 2, 3]), &chain, false) + .unwrap() ); // Still holds an old fragment: not caught up. assert_false!( - is_index_remap_caught_up(&version, &index_covering(5, &[1, 4, 5]), &chain).unwrap() + is_index_remap_caught_up(&version, &index_covering(5, &[1, 4, 5]), &chain, false) + .unwrap() ); // Bitmap advanced onto the new fragment but data not yet remapped: not // caught up (why the chain must include new frags). assert_false!( - is_index_remap_caught_up(&version, &index_covering(5, &[1, 6]), &chain).unwrap() + is_index_remap_caught_up(&version, &index_covering(5, &[1, 6]), &chain, false).unwrap() ); // Once remapped (version advanced): caught up. assert_true!( - is_index_remap_caught_up(&version, &index_covering(11, &[1, 6]), &chain).unwrap() + is_index_remap_caught_up(&version, &index_covering(11, &[1, 6]), &chain, false) + .unwrap() ); } @@ -345,7 +367,72 @@ mod tests { // Stale index (version 5) covering only v2's new fragment [7]: not // disjoint from the chain, so not caught up on v1. - assert_false!(is_index_remap_caught_up(&v1, &index_covering(5, &[1, 7]), &chain).unwrap()); + assert_false!( + is_index_remap_caught_up(&v1, &index_covering(5, &[1, 7]), &chain, false).unwrap() + ); + } + + /// Stable-id entries never need the mapping, so the version gate does not + /// apply; only coverage that still names a rewritten fragment retains it. + #[test] + fn test_caught_up_stable_domain_needs_only_coverage() { + let version = reuse_version(10, &[4, 5], &[6]); + let chain = reuse_chain_frag_bitmap(std::slice::from_ref(&version)); + + // Coverage moved at commit, data untouched by design: caught up despite + // version 5 < 10. + assert_true!( + is_index_remap_caught_up(&version, &index_covering(5, &[1, 6]), &chain, true).unwrap() + ); + // Committed concurrently against the rewritten fragments: retained until + // the coverage is repaired. + assert_false!( + is_index_remap_caught_up(&version, &index_covering(5, &[1, 4, 5]), &chain, true) + .unwrap() + ); + } + + /// An index this build cannot classify may store addresses, so only coverage + /// disjoint from the chain counts as caught up. + #[test] + fn test_unknown_index_domain_retains_reuse_version() { + let version = reuse_version(10, &[4, 5], &[6]); + let chain = reuse_chain_frag_bitmap(std::slice::from_ref(&version)); + let mut index = index_covering(10, &[1, 6]); + index.index_details = Some(std::sync::Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.FutureAddressIndexDetails".to_string(), + value: vec![], + })); + let is_data_domain_stable = + effective_identifier_domain(&index, true).unwrap() == IdentifierDomain::StableRowId; + assert_false!(is_data_domain_stable); + assert_false!( + is_index_remap_caught_up(&version, &index, &chain, is_data_domain_stable).unwrap() + ); + + index.fragment_bitmap = Some(RoaringBitmap::from_iter([1u32, 2])); + assert_true!( + is_index_remap_caught_up(&version, &index, &chain, is_data_domain_stable).unwrap() + ); + } + + /// A known type at a version this build cannot read may have changed domain. + #[test] + fn test_unsupported_index_version_retains_reuse_version() { + let version = reuse_version(10, &[4, 5], &[6]); + let chain = reuse_chain_frag_bitmap(std::slice::from_ref(&version)); + let mut index = index_covering(10, &[1, 6]); + index.index_details = Some(std::sync::Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.BTreeIndexDetails".to_string(), + value: vec![], + })); + index.index_version = 99; + let is_data_domain_stable = + effective_identifier_domain(&index, true).unwrap() == IdentifierDomain::StableRowId; + assert_false!(is_data_domain_stable); + assert_false!( + is_index_remap_caught_up(&version, &index, &chain, is_data_domain_stable).unwrap() + ); } /// Whole-fragment removal (every row deleted, no replacement): an index @@ -359,11 +446,13 @@ mod tests { let chain = reuse_chain_frag_bitmap(std::slice::from_ref(&version)); // Index emptied by the deletion (empty bitmap): caught up. - assert_true!(is_index_remap_caught_up(&version, &index_covering(5, &[]), &chain).unwrap()); + assert_true!( + is_index_remap_caught_up(&version, &index_covering(5, &[]), &chain, false).unwrap() + ); // Bitmap still lists the removed fragment (not yet updated): retained. assert_false!( - is_index_remap_caught_up(&version, &index_covering(5, &[7]), &chain).unwrap() + is_index_remap_caught_up(&version, &index_covering(5, &[7]), &chain, false).unwrap() ); } @@ -423,6 +512,7 @@ mod tests { &frag_reuse_details.versions[0], scalar_index, &reuse_chain_frag_bitmap(&frag_reuse_details.versions), + false, ) .unwrap() ); @@ -438,6 +528,7 @@ mod tests { &frag_reuse_details.versions[0], scalar_index, &reuse_chain_frag_bitmap(&frag_reuse_details.versions), + false, ) .unwrap() ); @@ -537,6 +628,7 @@ mod tests { &frag_reuse_details.versions[0], index, &reuse_chain_frag_bitmap(&frag_reuse_details.versions), + false, ) .unwrap(), "index {col}_idx was not caught up after remap" diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 9632d338035..7f1e5361379 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -80,6 +80,7 @@ //! the successful tasks can be committed. You can also commit in batches if //! you wish. As long as the tasks don't rewrite any of the same fragments, //! they can be committed in any order. +use lance_core::utils::address::RowAddress; use lance_core::utils::row_addr_remap::{GroupInput, RowAddrRemap}; use std::borrow::Cow; use std::collections::{HashMap, HashSet}; @@ -811,15 +812,6 @@ impl CompactionPlanner for DefaultCompactionPlanner { dataset.manifest.data_storage_format.lance_file_format(), write_version, )?; - if self.options.defer_index_remap && dataset.manifest.uses_stable_row_ids() { - return Err(Error::invalid_input( - "defer_index_remap=true is not supported on datasets with stable row IDs: \ - stable row IDs do not require index remapping during compaction, so there \ - is nothing to defer." - .to_string(), - )); - } - // get_fragments should be returning fragments in sorted order (by id) // and fragment ids should be unique let fragments = dataset.get_fragments(); @@ -830,7 +822,8 @@ impl CompactionPlanner for DefaultCompactionPlanner { ); // Without stable row ids a rewrite moves every row address, so the // indices over the rewritten fragments have to be remapped in the same - // commit. An index this build cannot open cannot be remapped, so its + // commit (unless deferred through the fragment reuse index). An index + // this build cannot open cannot be remapped, so its // fragments join the caller's own exclusions and are left uncompacted - // taking the same path, which terminates the current bin rather than // letting the candidates on either side of the gap be planned together. @@ -2179,7 +2172,8 @@ async fn load_index_fragmaps(dataset: &Dataset) -> Result> { // written by a newer Lance. The same bitmaps also decide, through // `any_group_indexed`, whether a deferred compaction writes the // fragment-reuse index that a build which can read that index needs to - // repair its coverage. + // repair its coverage, or that an index storing stable row ids needs only + // to follow the rewritten fragments. let indices = load_all_indices(dataset).await?; let mut index_fragmaps = Vec::with_capacity(indices.len()); // System indices (fragment-reuse, mem-wal) don't define data coverage and @@ -2304,14 +2298,14 @@ pub struct RewriteResult { pub read_version: u64, /// The original fragments being replaced pub original_fragments: Vec, - /// Serialized `RoaringTreemap` of the row addresses from the original - /// fragments that were read during compaction. + /// Serialized `RoaringTreemap` of the live row addresses of the original + /// fragments, in the order compaction read them. /// - /// - `None` when configured with stable row IDs because the row ID - /// sequences are rechunked directly. - /// - `Some` then these addresses are either (1) written to storage for - /// deferred index remap post-processing, or (2) used with reserved - /// fragment IDs to build old-to-new mappings. + /// - `Some` when the addresses may have a consumer: an index remapped in + /// this commit, or the fragment reuse index of a deferred remap. With + /// stable row ids the row id sequences are rechunked directly and the + /// addresses are derived from the fragment metadata. + /// - `None` otherwise. pub row_addrs: Option>, } @@ -2351,6 +2345,42 @@ async fn reserve_fragment_ids( Ok(()) } +/// The addresses of the live rows of `fragments`: every physical row minus the +/// deletion vector. Equal to what a scan of the fragments in order yields, which +/// is what the fragment reuse index records, without rereading the data. +async fn live_row_addrs(dataset: &Dataset, fragments: &[Fragment]) -> Result { + let mut addrs = RoaringTreemap::new(); + for fragment in fragments { + let fragment_id = u32::try_from(fragment.id).map_err(|_| { + Error::internal(format!( + "Fragment id {} is outside the row-address range", + fragment.id + )) + })?; + let physical_rows = fragment.physical_rows.ok_or_else(|| { + Error::internal(format!( + "Fragment {} is missing physical_rows after migration", + fragment.id + )) + })?; + let physical_rows = u32::try_from(physical_rows).map_err(|_| { + Error::internal(format!( + "Fragment {} has physical_rows={physical_rows} outside the row-address range", + fragment.id + )) + })?; + let start = u64::from(RowAddress::first_row(fragment_id)); + addrs.insert_range(start..start + u64::from(physical_rows)); + if let Some(deletion_file) = &fragment.deletion_file { + let deletions = read_dataset_deletion_file(dataset, fragment.id, deletion_file).await?; + for offset in deletions.iter() { + addrs.remove(u64::from(RowAddress::new_from_parts(fragment_id, offset))); + } + } + } + Ok(addrs) +} + /// Rewrite the files in a single task. /// /// This assumes that the dataset is the correct read version to be compacted. @@ -2391,13 +2421,17 @@ async fn rewrite_files( .iter() .map(|f| f.physical_rows.unwrap() as u64) .sum::(); - // 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()); + // Row addresses are only useful if something will consume them: an index + // to remap now, or a deferred remap through the FRI, which the commit may + // still opt into, so any index counts. With stable row ids the scan yields + // row ids, not addresses; those come from the fragment metadata instead + // (below). + let uses_stable_row_ids = dataset.manifest.uses_stable_row_ids(); + let row_addrs_have_consumer = options.defer_index_remap + || load_indices_for_remapping(dataset.as_ref()) + .await? + .is_some(); + let capture_row_addrs = !uses_stable_row_ids && row_addrs_have_consumer; let mut new_fragments: Vec; let task_id = uuid::Uuid::new_v4(); log::info!( @@ -2563,18 +2597,7 @@ async fn rewrite_files( if capture_row_addrs { let (tx, rx) = std::sync::mpsc::channel(); - let mut addrs = RoaringTreemap::new(); - for frag in &fragments { - let frag_id = frag.id as u32; - let count = u64::try_from(frag.physical_rows.unwrap_or(0)).map_err(|_| { - Error::internal(format!( - "Fragment {} has too many physical rows to represent as row addresses", - frag.id - )) - })?; - let start = u64::from(lance_core::utils::address::RowAddress::first_row(frag_id)); - addrs.insert_range(start..start + count); - } + let addrs = live_row_addrs(dataset.as_ref(), &fragments).await?; let captured = CapturedRowIds::AddressStyle(addrs); let _ = tx.send(captured); row_ids_rx = Some(rx); @@ -2600,31 +2623,34 @@ async fn rewrite_files( // Wrap in an async block so `?` returns into `row_addrs_result` and we can // run cleanup before propagating the error. let row_addrs_result: Result>> = async { - if let Some(row_ids_rx) = row_ids_rx { + if uses_stable_row_ids { + log::info!("Compaction task {}: rechunking stable row ids", task_id); + rechunk_stable_row_ids(dataset.as_ref(), &mut new_fragments, &fragments).await?; + recalc_versions_for_rewritten_fragments( + dataset.as_ref(), + &mut new_fragments, + &fragments, + ) + .await?; + } + let row_addrs = if let Some(row_ids_rx) = row_ids_rx { let captured_ids = row_ids_rx .try_recv() .map_err(|err| Error::internal(format!("Failed to receive row ids: {}", err)))?; - let mut row_addrs = captured_ids.row_addrs(None)?.into_owned(); - // Compaction reads whole fragments, so the captured addresses are - // dense per-fragment ranges; run containers (standard roaring - // format) shrink the persisted blob from O(rows) to O(runs) bytes. - row_addrs.optimize(); - let mut serialized = Vec::with_capacity(row_addrs.serialized_size()); - row_addrs.serialize_into(&mut serialized)?; - Ok(Some(serialized)) + captured_ids.row_addrs(None)?.into_owned() + } else if uses_stable_row_ids && row_addrs_have_consumer { + live_row_addrs(dataset.as_ref(), &fragments).await? } else { - if dataset.manifest.uses_stable_row_ids() { - log::info!("Compaction task {}: rechunking stable row ids", task_id); - rechunk_stable_row_ids(dataset.as_ref(), &mut new_fragments, &fragments).await?; - recalc_versions_for_rewritten_fragments( - dataset.as_ref(), - &mut new_fragments, - &fragments, - ) - .await?; - } - Ok(None) - } + return Ok(None); + }; + let mut row_addrs = row_addrs; + // Compaction reads whole fragments, so the captured addresses are + // dense per-fragment ranges; run containers (standard roaring + // format) shrink the persisted blob from O(rows) to O(runs) bytes. + row_addrs.optimize(); + let mut serialized = Vec::with_capacity(row_addrs.serialized_size()); + row_addrs.serialize_into(&mut serialized)?; + Ok(Some(serialized)) } .await; @@ -3050,7 +3076,8 @@ pub async fn commit_compaction( } else if !options.defer_index_remap && !has_address_style { // We need to reserve fragment ids here so that the fragment bitmap // can be updated for each index. Only needed for stable row IDs - // since address-style IDs were already reserved above. + // without a deferred remap, since address-style results (including a + // deferred remap with stable row ids) were already reserved above. let new_fragments = rewrite_groups .iter_mut() .flat_map(|group| group.new_fragments.iter_mut()) @@ -3135,25 +3162,30 @@ mod tests { use self::remapping::RemappedIndex; use super::*; use crate::dataset::WriteDestination; + use crate::dataset::builder::DatasetBuilder; use crate::dataset::index::frag_reuse::cleanup_frag_reuse_index; use crate::dataset::optimize::remapping::{transpose_row_addrs, transpose_row_ids_from_digest}; use crate::dataset::scanner::ColumnOrdering; use crate::index::frag_reuse::{load_frag_reuse_index_details, open_frag_reuse_index}; use crate::index::vector::{StageParams, VectorIndexParams}; + use crate::session::Session; use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount}; + use arrow_array::cast::AsArray; + use arrow_array::types::UInt64Type; use arrow_array::types::{Float32Type, Float64Type, Int32Type, Int64Type}; use arrow_array::{ ArrayRef, Float32Array, Int32Array, Int64Array, LargeBinaryArray, LargeStringArray, - PrimitiveArray, RecordBatch, RecordBatchIterator, + PrimitiveArray, RecordBatch, RecordBatchIterator, StringArray, }; use arrow_schema::{DataType, Field, Schema}; use arrow_select::concat::concat_batches; use async_trait::async_trait; - use lance_arrow::BLOB_META_KEY; + use lance_arrow::json::ARROW_JSON_EXT_NAME; + use lance_arrow::{ARROW_EXT_NAME_KEY, BLOB_META_KEY}; use lance_core::Error; - use lance_core::ROW_ID; use lance_core::utils::address::RowAddress; use lance_core::utils::tempfile::TempStrDir; + use lance_core::{ROW_ADDR, ROW_ID}; use lance_datagen::Dimension; use lance_file::version::LanceFileVersion; use lance_index::frag_reuse::CompactFragReuseIndexHandle; @@ -3165,8 +3197,11 @@ mod tests { use lance_index::vector::pq::PQBuildParams; use lance_index::{Index, IndexType}; use lance_linalg::distance::{DistanceType, MetricType}; + use lance_table::feature_flags::FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS; + use lance_table::format::IdentifierDomain; use lance_table::io::manifest::read_manifest_indexes; use lance_testing::datagen::{BatchGenerator, IncrementingInt32, RandomVector}; + use prost::Message; use rstest::rstest; use std::collections::HashSet; use std::io::Cursor; @@ -4653,6 +4688,7 @@ mod tests { #[case] indexed_column: &str, #[case] query: &str, #[case] expected_rows: usize, + #[values(false, true)] defer_index_remap: bool, ) { let mut dataset = lance_datagen::gen_batch() .col("i", lance_datagen::array::step::()) @@ -4683,9 +4719,16 @@ mod tests { .await .unwrap(); - compact_files(&mut dataset, CompactionOptions::default(), None) - .await - .unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + defer_index_remap, + ..Default::default() + }, + None, + ) + .await + .unwrap(); // The index only knows the pre-compaction fragments, so it must not claim to // cover the fragment they were rewritten into. @@ -4705,6 +4748,13 @@ mod tests { .is_none_or(|covered| covered.is_empty()), "compaction must not point an address-domain index at the fragments it wrote" ); + // Address-domain entries are what the FRI describes, so this index does + // get the FRI when one was recorded. + let frag_reuse_index = dataset + .frag_reuse_index_for(&index, &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(frag_reuse_index.is_some(), defer_index_remap); // Every fragment therefore falls back to a full scan, and the filter is answered // in full. @@ -4915,124 +4965,1245 @@ mod tests { assert_eq!(index.details.versions, source_details.versions); } + async fn non_system_index_uuids(dataset: &Dataset) -> HashSet { + dataset + .load_indices() + .await + .unwrap() + .iter() + .filter(|index| index.name != FRAG_REUSE_INDEX_NAME) + .map(|index| index.uuid) + .collect() + } + + async fn row_addrs_by_row_id(dataset: &Dataset) -> HashMap { + let mut scan = dataset.scan(); + scan.with_row_id().with_row_address(); + let batch = scan + .project::<&str>(&[]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .zip( + batch[ROW_ADDR] + .as_primitive::() + .values() + .iter() + .copied(), + ) + .collect() + } + + /// `(row id, i)` of the rows matching `filter`, sorted by row id. + async fn rows_matching(dataset: &Dataset, filter: &str) -> Vec<(u64, i32)> { + let mut scan = dataset.scan(); + scan.filter(filter).unwrap().with_row_id(); + let batch = scan + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + row_id_and_i(&batch) + } + + /// `(row id, i)` of the ten rows nearest the origin among the live rows of + /// fragment 0 (`60 <= i < 100`), sorted by row id. The prefilter keeps stable + /// ids the index still holds for deleted rows out of the top ten, and every + /// candidate id collides with a fragment 0 address. + async fn nearest_rows(dataset: &Dataset) -> Vec<(u64, i32)> { + let mut scan = dataset.scan(); + let query = Float32Array::from(vec![0.0f32; 16]); + scan.nearest("vec", &query, 10) + .unwrap() + .filter("i >= 60 AND i < 100") + .unwrap() + .prefilter(true) + .with_row_id(); + let batch = scan + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + row_id_and_i(&batch) + } + + fn row_id_and_i(batch: &RecordBatch) -> Vec<(u64, i32)> { + let mut rows = batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .zip( + batch["i"] + .as_primitive::() + .values() + .iter() + .copied(), + ) + .collect::>(); + rows.sort_unstable(); + rows + } + + async fn row_id_set(dataset: &Dataset) -> HashSet { + row_addrs_by_row_id(dataset).await.into_keys().collect() + } + + /// Binary copy synthesizes the FRI addresses from fragment metadata instead + /// of a scan; with stable row ids they must still be the dense physical ranges. #[tokio::test] - async fn test_defer_index_remap_rejected_with_stable_row_ids() { - let test_dir = TempStrDir::default(); - let test_uri = &test_dir; + async fn test_defer_index_remap_binary_copy_with_stable_row_ids() { + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(400), + "memory://test/binary_copy", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + let row_ids_before = row_id_set(&dataset).await; + let filter = "i < 10 OR i = 250"; + let rows_before = rows_matching(&dataset, filter).await; + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 400, + defer_index_remap: true, + compaction_mode: Some(CompactionMode::ForceBinaryCopy), + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 1); + + let fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + let groups = &fri.details.versions[0].groups; + assert_eq!(groups.len(), 1); + let changed = + RoaringTreemap::deserialize_from(Cursor::new(&groups[0].changed_row_addrs)).unwrap(); + let mut expected = RoaringTreemap::new(); + for fragment_id in 0..4u32 { + let start = u64::from(RowAddress::first_row(fragment_id)); + expected.insert_range(start..start + 100); + } + assert_eq!(changed, expected); + assert_eq!(row_id_set(&dataset).await, row_ids_before); + assert_eq!(rows_matching(&dataset, filter).await, rows_before); + } + /// Distributed compaction: tasks executed apart from the commit carry the + /// addresses a deferred remap records, with stable row ids too. + #[tokio::test] + async fn test_defer_index_remap_distributed_with_stable_row_ids() { let data = sample_data(); - let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 9000))], data.schema()); + let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 900))], data.schema()); let mut dataset = Dataset::write( reader, - test_uri, + "memory://test/distributed", Some(WriteParams { - max_rows_per_file: 1000, // 9 fragments enable_stable_row_ids: true, + max_rows_per_file: 100, ..Default::default() }), ) .await .unwrap(); - assert!(dataset.manifest.uses_stable_row_ids()); + create_scalar_index(&mut dataset, "a", false).await; + let row_ids_before = row_id_set(&dataset).await; let options = CompactionOptions { - target_rows_per_fragment: 3_000, + target_rows_per_fragment: 300, defer_index_remap: true, ..Default::default() }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks().len(), 3); + let dataset_ref = &dataset; + let results = futures::stream::iter(plan.compaction_tasks()) + .then(|task| async move { task.execute(dataset_ref).await.unwrap() }) + .collect::>() + .await; + for result in &results { + assert!(result.row_addrs.is_some()); + } - // Fails at planning time, before any fragment is rewritten. - let plan_err = plan_compaction(&dataset, &options).await.unwrap_err(); - assert!(matches!(plan_err, Error::InvalidInput { .. })); - let msg = plan_err.to_string(); - assert!(msg.contains("defer_index_remap")); - assert!(msg.contains("stable row IDs")); - - // The full compact_files entry point fails the same way and leaves the - // dataset untouched (no new manifest version, no orphaned data files). - let version_before = dataset.manifest.version; - let compact_err = compact_files(&mut dataset, options, None) - .await - .unwrap_err(); - assert!(matches!(compact_err, Error::InvalidInput { .. })); - assert_eq!(dataset.manifest.version, version_before); - } - + commit_compaction( + &mut dataset, + results, + Arc::new(IgnoreRemap::default()), + &options, + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + let fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + assert_eq!(fri.details.versions[0].groups.len(), 3); + assert_eq!(row_id_set(&dataset).await, row_ids_before); + } + + /// An index built against the pre-compaction version and committed after a + /// deferred rewrite persists the rewritten fragment in its coverage. + /// Stable-id entries need only that coverage repaired, which + /// `remap_column_index` does without touching the index files; an + /// address-domain index cannot be remapped and is rebuilt instead. On a + /// shallow clone the kept files stay in the source dataset, so the repaired + /// metadata must keep pointing there. + #[rstest] + #[case::btree(BuiltinIndexType::BTree, IndexType::BTree)] + #[case::zone_map(BuiltinIndexType::ZoneMap, IndexType::ZoneMap)] #[tokio::test] - async fn test_defer_index_remap() { + async fn test_concurrent_index_creation_with_deferred_stable_row_id_compaction( + #[case] builtin: BuiltinIndexType, + #[case] index_type: IndexType, + #[values(false, true)] on_clone: bool, + ) { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); let mut data_gen = BatchGenerator::new() - .col(Box::new( - RandomVector::new().vec_width(128).named("vec".to_owned()), - )) - .col(Box::new(IncrementingInt32::new().named("i".to_owned()))); - + .col(Box::new(IncrementingInt32::new().named("i".to_owned()))) + .col(Box::new(IncrementingInt32::new().named("j".to_owned()))); let mut dataset = Dataset::write( - data_gen.batch(6_000), - "memory://test/table", + data_gen.batch(500), + test_uri, Some(WriteParams { - max_rows_per_file: 1_000, // 6 files + enable_stable_row_ids: true, + max_rows_per_file: 100, ..Default::default() }), ) .await .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + async fn persisted_coverage(dataset: &Dataset) -> IndexMetadata { + read_manifest_indexes( + &dataset.object_store, + &dataset.manifest_location, + &dataset.manifest, + ) + .await + .unwrap() + .into_iter() + .find(|index| index.name == "j_idx") + .unwrap() + } + let params = ScalarIndexParams::for_builtin(builtin); + dataset + .create_index(&["j"], index_type, Some("j_idx".into()), ¶ms, false) + .await + .unwrap(); + let racing = persisted_coverage(&dataset).await; + assert!(racing.fragment_bitmap.as_ref().unwrap().contains(0)); - // Create another same dataset to mimic behavior without deferred index remap - let mut data_gen2 = BatchGenerator::new() - .col(Box::new( - RandomVector::new().vec_width(128).named("vec".to_owned()), - )) - .col(Box::new(IncrementingInt32::new().named("i".to_owned()))); - - let mut dataset2 = Dataset::write( - data_gen2.batch(6_000), - "memory://test/table", - Some(WriteParams { - max_rows_per_file: 1_000, // 6 files + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, ..Default::default() - }), + }, + None, ) .await .unwrap(); - // Delete some rows to create deletions - dataset.delete("i < 500").await.unwrap(); - dataset2.delete("i < 500").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(); - assert_eq!(initial_indices.len(), 1); - assert_eq!(initial_indices[0].name, "scalar"); - - // Store the original scalar index UUID for comparison - let original_scalar_uuid = initial_indices[0].uuid; + // What a CreateIndex built against fragment 0 and committed after the + // rewrite persists. Committed directly: the resolver's FRI-aware path + // needs the FRI in the Rewrite transaction file, which does not carry it. + let current = persisted_coverage(&dataset).await; + dataset + .apply_commit( + Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![racing.clone()], + removed_indices: vec![current], + }, + None, + ), + &Default::default(), + &Default::default(), + ) + .await + .unwrap(); - // Plan and execute compaction manually - let options = CompactionOptions { - target_rows_per_fragment: 2_000, - defer_index_remap: true, - ..Default::default() - }; - let options2 = CompactionOptions { - target_rows_per_fragment: 2_000, - defer_index_remap: false, - ..Default::default() + let clone_dir = TempStrDir::default(); + let uri = if on_clone { + let clone_uri = format!("{}/clone", clone_dir); + dataset + .shallow_clone(clone_uri.as_str(), dataset.manifest.version, None) + .await + .unwrap(); + clone_uri + } else { + test_uri.to_string() }; - let plan = plan_compaction(&dataset, &options).await.unwrap(); - let plan2 = plan_compaction(&dataset2, &options2).await.unwrap(); + let expected = (0..10) + .chain([70]) + .map(|value| (value as u64, value)) + .collect::>(); + let filter = "j < 10 OR j = 70"; + let mut dataset = DatasetBuilder::from_uri(&uri) + .with_session(Arc::new(Session::default())) + .load() + .await + .unwrap(); + let persisted = persisted_coverage(&dataset).await; + assert!(persisted.fragment_bitmap.as_ref().unwrap().contains(0)); + assert_eq!(persisted.base_id.is_some(), on_clone); + let loaded = dataset.load_index_by_name("j_idx").await.unwrap().unwrap(); + let loaded_coverage = loaded.fragment_bitmap.as_ref().unwrap(); + assert!(!loaded_coverage.contains(0)); + assert!(loaded_coverage.contains(5)); + assert_eq!(rows_matching(&dataset, filter).await, expected); - let mut expected_all_old_frag_ids = Vec::new(); - let mut expected_all_new_frag_ids = Vec::new(); - let mut expected_all_new_frag_bitmap = RoaringBitmap::new(); - let mut expected_all_row_id_map = HashMap::new(); - let mut deferred_results = Vec::new(); - let mut immediate_results = Vec::new(); + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + let fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + assert_eq!( + fri.details.versions.len(), + 1, + "stale coverage still needs the mapping" + ); + + let version_before_remap = dataset.manifest.version; + remapping::remap_column_index(&mut dataset, &["j"], Some("j_idx".into())) + .await + .unwrap(); + let repaired = persisted_coverage(&dataset).await; + match index_type { + IndexType::BTree => { + assert_eq!(dataset.manifest.version, version_before_remap + 1); + assert_eq!(repaired.uuid, persisted.uuid, "metadata-only repair"); + assert_eq!(repaired.files, persisted.files); + assert_eq!(repaired.base_id, persisted.base_id); + assert!(!repaired.fragment_bitmap.as_ref().unwrap().contains(0)); + assert!(repaired.fragment_bitmap.as_ref().unwrap().contains(5)); + assert!(repaired.dataset_version > persisted.dataset_version); + } + _ => { + // Address-domain entries cannot be remapped, so nothing is committed. + assert_eq!(dataset.manifest.version, version_before_remap); + assert_eq!(repaired, persisted); + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert_eq!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .details + .versions + .len(), + 1 + ); + dataset + .create_index(&["j"], index_type, Some("j_idx".into()), ¶ms, true) + .await + .unwrap(); + } + } + assert_eq!(rows_matching(&dataset, filter).await, expected); + + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .is_empty() + ); + let dataset = DatasetBuilder::from_uri(&uri) + .with_session(Arc::new(Session::default())) + .load() + .await + .unwrap(); + assert_eq!(rows_matching(&dataset, filter).await, expected); + } + + /// The FRI is not part of a stable-id index's cache identity, so a new FRI + /// version does not evict it. + #[tokio::test] + async fn test_stable_row_id_index_cache_survives_fri_change() { + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(500), + "memory://test/cache", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + let options = CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }; + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + compact_files(&mut dataset, options.clone(), None) + .await + .unwrap(); + let first_fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + let index_uuid = dataset + .load_index_by_name("scalar") + .await + .unwrap() + .unwrap() + .uuid; + dataset + .open_scalar_index("i", &index_uuid, &NoOpMetricsCollector) + .await + .unwrap(); + + dataset.delete("i >= 110 AND i < 160").await.unwrap(); + compact_files(&mut dataset, options, None).await.unwrap(); + let second_fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + assert_ne!(first_fri.uuid, second_fri.uuid); + + // Same cache namespace as before: served from cache, no new entry. + let session = dataset.session(); + let stats_before = session.index_cache_stats().await; + dataset + .open_scalar_index("i", &index_uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let stats_after = session.index_cache_stats().await; + assert_eq!(stats_after.num_entries, stats_before.num_entries); + assert!(stats_after.hits > stats_before.hits); + } + + /// Merging new data into indices after a deferred compaction: the btree and + /// vector merge paths read the old segments, whose entries are stable row ids + /// colliding with fragment 0's addresses, and NGram must take the merge path + /// even though the FRI version postdates its segment. + #[tokio::test] + async fn test_optimize_indices_after_deferred_compaction_with_stable_row_ids() { + use lance_datagen::{BatchCount, RowCount}; + use lance_index::optimize::OptimizeOptions; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let batch_reader = |start: i32, rows: u64| { + lance_datagen::gen_batch() + .col( + "i", + lance_datagen::array::step_custom::(start, 1), + ) + .col( + "s", + lance_datagen::array::cycle_utf8_literals(&["needle", "haystack"]), + ) + .col( + "vec", + lance_datagen::array::rand_vec::(Dimension::from(16)), + ) + .into_reader_rows(RowCount::from(rows), BatchCount::from(1)) + }; + let mut dataset = Dataset::write( + batch_reader(0, 500), + test_uri, + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + dataset + .create_index( + &["s"], + IndexType::NGram, + Some("s_idx".into()), + &ScalarIndexParams::for_builtin(BuiltinIndexType::NGram), + false, + ) + .await + .unwrap(); + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("vector".into()), + &VectorIndexParams::ivf_pq(1, 8, 1, MetricType::L2, 50), + false, + ) + .await + .unwrap(); + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + let fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + let ngram = dataset.load_index_by_name("s_idx").await.unwrap().unwrap(); + assert!( + crate::index::append::fragment_reuse_affects_segments(&fri, [&ngram]), + "the FRI version postdates the segment, so an ungated merge would rebuild" + ); + + let dataset = Dataset::write( + batch_reader(500, 100), + test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + let mut dataset = dataset; + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri) + .with_session(Arc::new(Session::default())) + .load() + .await + .unwrap(); + let mut coverage_by_index: HashMap = HashMap::new(); + for index in dataset.load_indices().await.unwrap().iter() { + if index.name != FRAG_REUSE_INDEX_NAME { + *coverage_by_index.entry(index.name.clone()).or_default() |= + index.fragment_bitmap.as_ref().unwrap(); + } + } + let live_fragments: RoaringBitmap = + dataset.fragments().iter().map(|f| f.id as u32).collect(); + assert_eq!(live_fragments.len(), 6); + for (name, coverage) in &coverage_by_index { + assert_eq!( + coverage, &live_fragments, + "{name} covers every fragment after optimize" + ); + } + let filter = "i < 10 OR i = 70 OR i >= 590"; + let expected = (0..10) + .chain([70]) + .chain(590..600) + .map(|value| (value as u64, value)) + .collect::>(); + assert_eq!(rows_matching(&dataset, filter).await, expected); + + let mut with_index = dataset.scan(); + with_index + .filter("contains(s, 'needle')") + .unwrap() + .with_row_id(); + let with_index = with_index + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut without_index = dataset.scan(); + without_index + .filter("contains(s, 'needle')") + .unwrap() + .with_row_id() + .use_scalar_index(false); + let without_index = without_index + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(with_index.num_rows(), 275); + assert_eq!(row_id_and_i(&with_index), row_id_and_i(&without_index)); + + let live_row_ids = row_id_set(&dataset).await; + let nearest = nearest_rows(&dataset).await; + assert_eq!(nearest.len(), 10); + for (row_id, value) in nearest { + assert!(live_row_ids.contains(&row_id), "row id {row_id}"); + assert_eq!(row_id, value as u64); + } + } + + /// A JSON index answers in its target's domain. Over a zone map it stores row + /// addresses, so on a stable-row-id dataset it loses the rewritten fragments + /// and receives the FRI; over a btree it stores stable row ids and does neither. + #[rstest] + #[case::zone_map("zonemap", true)] + #[case::btree("btree", false)] + #[tokio::test] + async fn test_json_index_domain_with_stable_row_ids( + #[case] target_index_type: &str, + #[case] stores_row_addrs: bool, + ) { + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("json", DataType::Utf8, false).with_metadata(metadata), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..200)), + Arc::new(StringArray::from_iter_values( + (0..200).map(|value| format!(r#"{{"val": {value}}}"#)), + )), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://test/json", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + let params = ScalarIndexParams::new("json".to_string()).with_params(&serde_json::json!({ + "target_index_type": target_index_type, + "path": "val", + })); + dataset + .create_index( + &["json"], + IndexType::Scalar, + Some("json_idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + // Collides stable ids 0..100 with fragment 0's addresses. + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + let filter = "json_get_int(json, 'val') < 10 OR json_get_int(json, 'val') = 70"; + let rows_before = rows_matching(&dataset, filter).await; + assert_eq!(rows_before.len(), 11); + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + let index = dataset + .load_index_by_name("json_idx") + .await + .unwrap() + .unwrap(); + assert_eq!( + index.identifier_domain(true).unwrap() == IdentifierDomain::RowAddress, + stores_row_addrs + ); + assert_eq!( + index.fragment_bitmap.as_ref().unwrap().contains(2), + !stores_row_addrs + ); + assert_eq!( + dataset + .frag_reuse_index_for(&index, &NoOpMetricsCollector) + .await + .unwrap() + .is_some(), + stores_row_addrs + ); + assert_eq!(rows_matching(&dataset, filter).await, rows_before); + } + + fn unknown_domain_index( + field_id: i32, + dataset_version: u64, + covered: &[u32], + details: prost_types::Any, + ) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + name: "future_idx".to_string(), + fields: vec![field_id], + covering_fields: vec![], + dataset_version, + fragment_bitmap: Some(RoaringBitmap::from_iter(covered.iter().copied())), + index_details: Some(Arc::new(details)), + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + fn future_index_details() -> prost_types::Any { + prost_types::Any { + type_url: "type.googleapis.com/lance.table.FutureAddressIndexDetails".to_string(), + value: vec![], + } + } + + async fn commit_index(dataset: &mut Dataset, index: IndexMetadata) { + dataset + .apply_commit( + Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![index], + removed_indices: vec![], + }, + None, + ), + &Default::default(), + &Default::default(), + ) + .await + .unwrap(); + } + + /// An index whose type this build cannot classify may store row addresses, so a + /// rewrite withdraws the rewritten fragments from its coverage, deferred or not. + #[rstest] + #[case::eager(false)] + #[case::deferred(true)] + #[tokio::test] + async fn test_unknown_index_domain_loses_rewritten_coverage( + #[case] defer_index_remap: bool, + #[values(false, true)] unsupported_version: bool, + ) { + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(500), + "memory://test/unknown_domain", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + let field_id = dataset.schema().field("i").unwrap().id; + let version = dataset.manifest.version; + let mut index = + unknown_domain_index(field_id, version, &[0, 1, 2, 3, 4], future_index_details()); + if unsupported_version { + index.index_details = Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.BTreeIndexDetails".to_string(), + value: vec![], + })); + index.index_version = 99; + } + commit_index(&mut dataset, index).await; + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(dataset.get_fragment(5).is_some()); + + let persisted = read_manifest_indexes( + &dataset.object_store, + &dataset.manifest_location, + &dataset.manifest, + ) + .await + .unwrap() + .into_iter() + .find(|index| index.name == "future_idx") + .unwrap(); + assert_eq!( + persisted.fragment_bitmap.unwrap(), + RoaringBitmap::from_iter([1u32, 2, 3, 4]) + ); + assert_eq!( + dataset.frag_reuse_index().await.unwrap().is_some(), + defer_index_remap + ); + if defer_index_remap { + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .is_empty(), + "withdrawn coverage needs no mapping" + ); + } + assert_eq!(rows_matching(&dataset, "i < 10 OR i = 70").await.len(), 11); + } + + /// An index built from the version a deferred compaction read shares the FRI + /// version's watermark, and the rebased rewrite advances its coverage while its + /// entries still hold old addresses. Cleanup must keep the version and remap + /// must rewrite the index. + #[tokio::test] + async fn test_equal_watermark_index_is_retained_and_remapped() { + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(500), + "memory://test/equal_watermark", + Some(WriteParams { + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + // An existing index makes the compaction record an FRI. + create_scalar_index(&mut dataset, "i", false).await; + let mut stale = dataset.clone(); + dataset + .create_index( + &["i"], + IndexType::Scalar, + Some("second".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + compact_files( + &mut stale, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + dataset.drop_index("scalar").await.unwrap(); + + let watermark = dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .details + .versions[0] + .dataset_version; + let second = read_manifest_indexes( + &dataset.object_store, + &dataset.manifest_location, + &dataset.manifest, + ) + .await + .unwrap() + .into_iter() + .find(|index| index.name == "second") + .unwrap(); + assert_eq!(second.dataset_version, watermark); + let coverage = second.fragment_bitmap.as_ref().unwrap(); + assert!(coverage.contains(5) && !coverage.contains(0)); + + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert_eq!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .details + .versions + .len(), + 1 + ); + remapping::remap_column_index(&mut dataset, &["i"], Some("second".into())) + .await + .unwrap(); + let remapped = dataset.load_index_by_name("second").await.unwrap().unwrap(); + assert_ne!(remapped.uuid, second.uuid); + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .is_empty() + ); + assert_eq!(rows_matching(&dataset, "i < 10 OR i = 70").await.len(), 11); + } + + /// A JSON index over a target this build cannot read is unusable, and cleanup + /// keeps the mapping while its coverage still names a rewritten fragment. + #[tokio::test] + async fn test_json_index_over_unknown_target_is_refused() { + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(500), + "memory://test/json_unknown_target", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + // Committed after the rewrite with coverage of the rewritten fragment 0. + let field_id = dataset.schema().field("i").unwrap().id; + let json_details = prost_types::Any { + type_url: "type.googleapis.com/lance.index.pb.JsonIndexDetails".to_string(), + value: lance_index::pb::JsonIndexDetails { + path: "val".to_string(), + target_details: Some(future_index_details()), + } + .encode_to_vec(), + }; + let version = dataset.manifest.version; + commit_index( + &mut dataset, + unknown_domain_index(field_id, version, &[0, 1, 2, 3, 4], json_details), + ) + .await; + + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + assert_eq!( + dataset + .frag_reuse_index() + .await + .unwrap() + .unwrap() + .details + .versions + .len(), + 1, + "coverage of a rewritten fragment keeps the mapping" + ); + + // Unusable to this build: planning ignores it instead of panicking on the target. + assert_eq!(rows_matching(&dataset, "i < 10 OR i = 70").await.len(), 11); + assert!( + !dataset + .load_indices() + .await + .unwrap() + .iter() + .any(|index| index.name == "future_idx") + ); + assert!( + remapping::remap_column_index(&mut dataset, &["i"], Some("future_idx".into())) + .await + .is_err() + ); + + let index = read_manifest_indexes( + &dataset.object_store, + &dataset.manifest_location, + &dataset.manifest, + ) + .await + .unwrap() + .into_iter() + .find(|index| index.name == "future_idx") + .unwrap(); + assert_eq!( + index.identifier_domain(true).unwrap(), + IdentifierDomain::Unknown + ); + let load_err = dataset + .frag_reuse_index_for(&index, &NoOpMetricsCollector) + .await + .unwrap_err(); + assert!(matches!(load_err, Error::NotSupported { .. }), "{load_err}"); + } + + /// Deferred remap on a stable-row-id dataset records an FRI over physical + /// addresses and leaves the indices, whose entries are stable row ids, alone. + /// Fragment 0 is the rewritten group, so its addresses `0..100` collide with + /// stable ids `0..100`: an index wrongly remapped through the FRI would + /// answer with addresses in the new fragment instead of those ids. + #[tokio::test] + async fn test_defer_index_remap_with_stable_row_ids() { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let mut data_gen = BatchGenerator::new() + .col(Box::new( + RandomVector::new().vec_width(16).named("vec".to_owned()), + )) + .col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + let mut dataset = Dataset::write( + data_gen.batch(500), + test_uri, + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); + create_scalar_index(&mut dataset, "i", false).await; + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("vector".into()), + &VectorIndexParams::ivf_pq(1, 8, 1, MetricType::L2, 50), + false, + ) + .await + .unwrap(); + + // Half of fragment 0: it becomes the only compaction candidate, the other + // fragments already hold `target_rows_per_fragment` rows. + dataset.delete("i >= 10 AND i < 60").await.unwrap(); + + let indices_before = non_system_index_uuids(&dataset).await; + let filter = "i < 10 OR i = 70 OR i = 150"; + let scalar_before = rows_matching(&dataset, filter).await; + assert_eq!(scalar_before.len(), 12); + let nearest_before = nearest_rows(&dataset).await; + let addrs_before = row_addrs_by_row_id(&dataset).await; + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + // Cold reopen: persisted metadata and empty index caches. + let mut dataset = DatasetBuilder::from_uri(test_uri) + .with_session(Arc::new(Session::default())) + .load() + .await + .unwrap(); + assert!(dataset.get_fragment(0).is_none()); + assert!(dataset.get_fragment(5).is_some()); + let fri = dataset + .frag_reuse_index() + .await + .unwrap() + .expect("deferred compaction records an FRI"); + assert_eq!(fri.details.versions.len(), 1); + assert_eq!(non_system_index_uuids(&dataset).await, indices_before); + for index in dataset + .load_indices() + .await + .unwrap() + .iter() + .filter(|index| index.name != FRAG_REUSE_INDEX_NAME) + { + let coverage = index.fragment_bitmap.as_ref().unwrap(); + assert!(!coverage.contains(0), "{}: {coverage:?}", index.name); + assert!(coverage.contains(5), "{}: {coverage:?}", index.name); + } + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0 + ); + + assert_eq!(rows_matching(&dataset, filter).await, scalar_before); + assert_eq!(nearest_rows(&dataset).await, nearest_before); + + let addrs_after = row_addrs_by_row_id(&dataset).await; + let remap = fri.row_addr_remap(); + for row_id in [0u64, 9, 60, 99] { + assert_eq!( + remap.get(addrs_before[&row_id]), + Some(Some(addrs_after[&row_id])), + "row id {row_id}" + ); + } + assert_eq!( + remap.get(u64::from(RowAddress::new_from_parts(0, 10))), + Some(None), + "deleted before the rewrite" + ); + assert_eq!(remap.get(addrs_before[&150]), None, "untouched fragment"); + assert_eq!(addrs_before[&150], addrs_after[&150]); + + // A shallow clone copies the fragments and the FRI without entering + // `build_manifest`, so the fence has to survive `Manifest::shallow_clone`. + let clone_dir = TempStrDir::default(); + let clone_uri = format!("{}/clone", clone_dir); + let cloned = dataset + .shallow_clone(clone_uri.as_str(), dataset.manifest.version, None) + .await + .unwrap(); + assert!(cloned.manifest.uses_stable_row_ids()); + assert!( + cloned + .load_index_by_name(FRAG_REUSE_INDEX_NAME) + .await + .unwrap() + .is_some(), + "precondition: the clone carries the FRI" + ); + assert_ne!( + cloned.manifest.reader_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0, + "a clone of a stable-row-id table with an FRI must stay fenced for readers" + ); + assert_ne!( + cloned.manifest.writer_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0, + "a clone of a stable-row-id table with an FRI must stay fenced for writers" + ); + + // Every index stores stable row ids, so the version drains at once. + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + let fri = dataset.frag_reuse_index().await.unwrap().unwrap(); + assert!(fri.is_empty()); + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS, + 0, + "an emptied FRI still fences older readers" + ); + } + + #[tokio::test] + async fn test_defer_index_remap() { + let mut data_gen = BatchGenerator::new() + .col(Box::new( + RandomVector::new().vec_width(128).named("vec".to_owned()), + )) + .col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + + let mut dataset = Dataset::write( + data_gen.batch(6_000), + "memory://test/table", + Some(WriteParams { + max_rows_per_file: 1_000, // 6 files + ..Default::default() + }), + ) + .await + .unwrap(); + + // Create another same dataset to mimic behavior without deferred index remap + let mut data_gen2 = BatchGenerator::new() + .col(Box::new( + RandomVector::new().vec_width(128).named("vec".to_owned()), + )) + .col(Box::new(IncrementingInt32::new().named("i".to_owned()))); + + let mut dataset2 = Dataset::write( + data_gen2.batch(6_000), + "memory://test/table", + Some(WriteParams { + max_rows_per_file: 1_000, // 6 files + ..Default::default() + }), + ) + .await + .unwrap(); + + // Delete some rows to create deletions + dataset.delete("i < 500").await.unwrap(); + dataset2.delete("i < 500").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(); + assert_eq!(initial_indices.len(), 1); + assert_eq!(initial_indices[0].name, "scalar"); + + // Store the original scalar index UUID for comparison + let original_scalar_uuid = initial_indices[0].uuid; + + // Plan and execute compaction manually + let options = CompactionOptions { + target_rows_per_fragment: 2_000, + defer_index_remap: true, + ..Default::default() + }; + let options2 = CompactionOptions { + target_rows_per_fragment: 2_000, + defer_index_remap: false, + ..Default::default() + }; + + let plan = plan_compaction(&dataset, &options).await.unwrap(); + let plan2 = plan_compaction(&dataset2, &options2).await.unwrap(); + + let mut expected_all_old_frag_ids = Vec::new(); + let mut expected_all_new_frag_ids = Vec::new(); + let mut expected_all_new_frag_bitmap = RoaringBitmap::new(); + let mut expected_all_row_id_map = HashMap::new(); + let mut deferred_results = Vec::new(); + let mut immediate_results = Vec::new(); for (task, task2) in plan.tasks().iter().zip(plan2.tasks()) { let deferred_result = rewrite_files(Cow::Borrowed(&dataset), task.clone(), &options) @@ -5196,8 +6367,11 @@ mod tests { assert_eq!(current_scalar_index.uuid, original_scalar_uuid); } + #[rstest] #[tokio::test] - async fn test_defer_index_remap_skips_fri_when_no_indexed_data() { + async fn test_defer_index_remap_skips_fri_when_no_indexed_data( + #[values(false, true)] enable_stable_row_ids: bool, + ) { // A deferred compaction touching no indexed data must write no FRI -- // such a version is un-drainable (remap no-ops, trim retains it forever). let mut data_gen = @@ -5208,6 +6382,7 @@ mod tests { "memory://test/noindex", Some(WriteParams { max_rows_per_file: 100, // 6 small files -> compaction has work + enable_stable_row_ids, ..Default::default() }), ) @@ -5242,8 +6417,11 @@ mod tests { ); } + #[rstest] #[tokio::test] - async fn test_defer_index_remap_multiple_compactions() { + async fn test_defer_index_remap_multiple_compactions( + #[values(false, true)] enable_stable_row_ids: bool, + ) { let mut data_gen = BatchGenerator::new() .col(Box::new( RandomVector::new().vec_width(128).named("vec".to_owned()), @@ -5255,6 +6433,7 @@ mod tests { "memory://test/table", Some(WriteParams { max_rows_per_file: 1_000, // 6 files + enable_stable_row_ids, ..Default::default() }), ) @@ -5406,8 +6585,11 @@ mod tests { } } + #[rstest] #[tokio::test] - async fn test_deferred_compaction_not_split_by_frag_reuse_index() { + async fn test_deferred_compaction_not_split_by_frag_reuse_index( + #[values(false, true)] enable_stable_row_ids: bool, + ) { // The fragment-reuse index is a system index and must be excluded from // compaction bin planning; otherwise its covered fragment is isolated and // the small fragments never coalesce back to one. @@ -5427,6 +6609,7 @@ mod tests { test_uri, Some(WriteParams { max_rows_per_file: 200, + enable_stable_row_ids, ..Default::default() }), ) diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 491ad893063..32d3edfda07 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -6,15 +6,15 @@ use crate::Result; use crate::dataset::transaction::{Operation, Transaction}; -use crate::index::DatasetIndexExt; use crate::index::frag_reuse::{load_frag_reuse_index_details, open_frag_reuse_index}; +use crate::index::{DatasetIndexExt, effective_identifier_domain}; use crate::{Dataset, index}; use async_trait::async_trait; use lance_core::Error; use lance_core::utils::address::RowAddress; use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragDigest}; -use lance_table::format::{Fragment, IndexFile, IndexMetadata}; +use lance_table::format::{Fragment, IdentifierDomain, IndexFile, IndexMetadata}; use lance_table::io::manifest::read_manifest_indexes; use roaring::RoaringTreemap; use serde::{Deserialize, Serialize}; @@ -245,11 +245,27 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { // remapping a *sibling* index). let baseline_version = curr_index_meta.dataset_version; let has_unknown_coverage = curr_index_meta.fragment_bitmap.is_none(); + // Stable row ids survive a rewrite, so such an index never needs its data + // remapped; only coverage committed against rewritten fragments (an index + // created concurrently with the compaction) can be stale. + let is_data_domain_stable = match effective_identifier_domain( + &curr_index_meta, + dataset.manifest.uses_stable_row_ids(), + )? { + IdentifierDomain::StableRowId => true, + IdentifierDomain::RowAddress => false, + IdentifierDomain::Unknown => { + return Err(Error::not_supported(format!( + "cannot remap index {} ({}): its identifier domain is unknown to this build", + curr_index_meta.name, curr_index_meta.uuid + ))); + } + }; let (should_remap, mut bitmap_after_remap) = match curr_index_meta.fragment_bitmap.clone() { Some(mut index_frag_bitmap) => { let mut should_remap = false; for version in frag_reuse_index.details.versions.iter() { - let data_predates_version = baseline_version < version.dataset_version; + let data_predates_version = baseline_version <= version.dataset_version; for group in version.groups.iter() { let mut old_frag_in_index = 0; for old_frag in group.old_frags.iter() { @@ -271,7 +287,8 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { } index_frag_bitmap.extend(group.new_frags.iter().map(|f| f.id as u32)); should_remap = true; - } else if data_predates_version + } else if !is_data_domain_stable + && data_predates_version && group .new_frags .iter() @@ -299,8 +316,11 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { // leaves missing mappings unchanged, so no composed per-row map is needed. // This also handles the sibling-coverage-remap case: remapping is driven by // the row addresses stored in the index, not by its already-advanced bitmap. - let remap_result = - index::remap_index(dataset, index_id, frag_reuse_index.row_addr_remap()).await?; + let remap_result = if is_data_domain_stable { + RemapResult::Keep(*index_id) + } else { + index::remap_index(dataset, index_id, frag_reuse_index.row_addr_remap()).await? + }; // Remapping advances the index watermark for fragment-reuse cleanup, but it // does not incorporate overlays committed after the source index was built. @@ -327,16 +347,18 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { }; let new_index_meta = match remap_result { - // Nothing to commit: either the composed remap emptied the index (every - // row deleted), matching the prior per-version behavior, or - // `index::remap_index` withdrew a covered index it cannot carry payload - // through. Either way the existing entry is left untouched. + // Nothing to commit: `index::remap_index` withdrew a covered index it + // cannot carry payload through, so the existing entry is left untouched. // - // The withdrawal case is unreachable here today: the only caller is - // `remap_column_index`, which refuses a covered index first. Compaction - // reaches that withdrawal through `DatasetIndexRemapper`, which handles - // `RemapResult::Drop` in `dataset/index.rs` rather than through here. + // Unreachable here today: the only caller is `remap_column_index`, which + // refuses a covered index first. Compaction reaches that withdrawal + // through `DatasetIndexRemapper`, which handles `RemapResult::Drop` in + // `dataset/index.rs` rather than through here. RemapResult::Drop => return Ok(()), + // Same files, new coverage: the composed remap emptied the index (every + // row deleted) or the index stores stable row ids and only its coverage + // was stale. The files stay where they are, which on a shallow clone is + // the source dataset, so the base travels with them. RemapResult::Keep(new_id) => IndexMetadata { uuid: new_id, name: curr_index_meta.name.clone(), @@ -347,7 +369,7 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { index_details: curr_index_meta.index_details.clone(), index_version: curr_index_meta.index_version, created_at: curr_index_meta.created_at, - base_id: None, + base_id: curr_index_meta.base_id, files: curr_index_meta.files.clone(), }, RemapResult::Remapped(remapped_index) => IndexMetadata { diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 415958a438a..dce50ed29b9 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -62,7 +62,9 @@ use lance_io::utils::{ read_version, }; use lance_table::format::{DataFile, Fragment, SelfDescribingFileReader}; -use lance_table::format::{IndexFile, IndexMetadata, list_index_files_with_sizes}; +use lance_table::format::{ + IdentifierDomain, IndexFile, IndexMetadata, list_index_files_with_sizes, +}; use lance_table::io::manifest::read_manifest_indexes; use roaring::RoaringBitmap; use scalar::index_matches_criteria; @@ -2061,7 +2063,9 @@ impl DatasetIndexExt for Dataset { let has_retired_coverage = segments .iter() .any(|segment| !(segment.fragment_bitmap() - &dataset_fragments).is_empty()); - let frag_reuse_index = self.open_frag_reuse_index(&NoOpMetricsCollector).await?; + let frag_reuse_index = self + .frag_reuse_index_for_row_id_entries(&NoOpMetricsCollector) + .await?; let requires_rebuild = frag_reuse_index.as_ref().is_some_and(|frag_reuse_index| { segments.iter().any(|segment| { append::fragment_reuse_affects_segment( @@ -2725,6 +2729,27 @@ async fn gather_fragment_statistics( ))) } +/// [`IndexMetadata::identifier_domain`] that also treats an index version this build +/// cannot read as unknown. +pub(crate) fn effective_identifier_domain( + index: &IndexMetadata, + uses_stable_row_ids: bool, +) -> Result { + if !uses_stable_row_ids { + return Ok(IdentifierDomain::RowAddress); + } + if unsupported_index_version(index).is_some() { + return Ok(IdentifierDomain::Unknown); + } + index.declared_identifier_domain() +} + +/// [`effective_identifier_domain`] for the manifest build, which classifies for +/// stable-row-id tables. +pub(crate) fn stable_row_id_index_domain(index: &IndexMetadata) -> Result { + effective_identifier_domain(index, true) +} + /// `None` when this build supports the index's version, otherwise the highest /// version it does support. /// @@ -2910,12 +2935,31 @@ pub trait DatasetIndexInternalExt: DatasetIndexExt { name: &str, ) -> Result; - /// Opens the fragment reuse index + /// Opens the fragment reuse index. Use it for fragment coverage; index data + /// goes through [`Self::frag_reuse_index_for`], which honors the index's + /// identifier domain. async fn open_frag_reuse_index( &self, metrics: &dyn MetricsCollector, ) -> Result>>; + /// The fragment reuse index to remap `index`'s stored identifiers through: `None` + /// when there is none or `index` stores stable row ids, an error when its domain is + /// unknown to this build. + async fn frag_reuse_index_for( + &self, + index: &IndexMetadata, + metrics: &dyn MetricsCollector, + ) -> Result>>; + + /// [`Self::frag_reuse_index_for`] for an index being built or merged, whose + /// entries are the `_rowid` column and which has no metadata to classify yet: + /// stable row ids on this dataset, row addresses otherwise. + async fn frag_reuse_index_for_row_id_entries( + &self, + metrics: &dyn MetricsCollector, + ) -> Result>>; + /// Opens the MemWAL index async fn open_mem_wal_index( &self, @@ -2953,7 +2997,14 @@ impl DatasetIndexInternalExt for Dataset { // Checking for cache existence is cheap so we just check the vector caches. // Scalar indices cache themselves inside `open_scalar_index` (the cache // key is a plugin detail), so there is no cheap scalar check here. - let frag_reuse_uuid = self.frag_reuse_index_uuid().await; + let index_meta = self + .load_index(uuid) + .await? + .ok_or_else(|| Error::index(format!("Index with id {} does not exist", uuid)))?; + let frag_reuse_uuid = self + .frag_reuse_index_for(&index_meta, metrics) + .await? + .map(|index| index.uuid); // Check sized cache for IvfIndexState (v2+ indices). let state_key = IvfIndexStateCacheKey::new(uuid, frag_reuse_uuid.as_ref()); @@ -2980,10 +3031,6 @@ impl DatasetIndexInternalExt for Dataset { // We determine if this is a vector index by checking if INDEX_FILE_NAME exists in the // file list (available since file sizes tracking was added). If the file list is not // available (older indices), we fall back to checking file existence via HEAD request. - let index_meta = self - .load_index(uuid) - .await? - .ok_or_else(|| Error::index(format!("Index with id {} does not exist", uuid)))?; // Check if this is a vector index by looking at the files list let is_vector_index = if let Some(files) = &index_meta.files { @@ -3032,11 +3079,12 @@ impl DatasetIndexInternalExt for Dataset { uuid: &Uuid, metrics: &dyn MetricsCollector, ) -> Result> { - let frag_reuse_uuid = self.frag_reuse_index_uuid().await; let index_meta = self .load_index(uuid) .await? .ok_or_else(|| Error::index(format!("Index with id {} does not exist", uuid)))?; + let frag_reuse_index = self.frag_reuse_index_for(&index_meta, metrics).await?; + let frag_reuse_uuid = frag_reuse_index.as_ref().map(|index| index.uuid); let object_store = self.object_store_for_index(&index_meta).await?; // Check sized cache first (v2+ indices with serializable state). @@ -3044,7 +3092,6 @@ impl DatasetIndexInternalExt for Dataset { if let Some(entry) = self.index_cache.get_with_key(&state_key).await { log::debug!("Found IvfIndexState in cache uuid: {}", uuid); let partition_cache = self.index_cache.for_index(uuid, frag_reuse_uuid.as_ref()); - let frag_reuse_index = self.open_frag_reuse_index(metrics).await?; return entry .0 .reconstruct( @@ -3062,7 +3109,6 @@ impl DatasetIndexInternalExt for Dataset { return Ok(cached.0.clone()); } - let frag_reuse_index = self.open_frag_reuse_index(metrics).await?; let index_dir = self.indice_files_dir(&index_meta)?; let index_file = index_dir .clone() @@ -3404,6 +3450,37 @@ impl DatasetIndexInternalExt for Dataset { } } + async fn frag_reuse_index_for( + &self, + index: &IndexMetadata, + metrics: &dyn MetricsCollector, + ) -> Result>> { + match effective_identifier_domain(index, self.manifest.uses_stable_row_ids())? { + IdentifierDomain::RowAddress => self.open_frag_reuse_index(metrics).await, + IdentifierDomain::StableRowId => Ok(None), + IdentifierDomain::Unknown => Err(Error::not_supported(format!( + "index {} ({}) stores identifiers of a type unknown to this build: {}", + index.name, + index.uuid, + index + .index_details + .as_ref() + .map_or("no details", |details| details.type_url.as_str()) + ))), + } + } + + async fn frag_reuse_index_for_row_id_entries( + &self, + metrics: &dyn MetricsCollector, + ) -> Result>> { + if self.manifest.uses_stable_row_ids() { + Ok(None) + } else { + self.open_frag_reuse_index(metrics).await + } + } + async fn open_mem_wal_index( &self, metrics: &dyn MetricsCollector, @@ -12330,6 +12407,11 @@ mod tests { /// Both sides therefore have to count the same indices: planning from the /// filtered view while the commit carries the complete one makes compaction /// fail outright on a dataset holding an index from a newer build. + /// + /// The index survives, but this build cannot tell whether a version it cannot + /// read still stores stable row ids, so the rewrite withdraws its coverage of + /// the rewritten fragments instead of advancing it. Queries fall back to + /// scanning until a writer that can read the index rebuilds that coverage. #[tokio::test] async fn test_compaction_survives_an_unsupported_index() { let test_dir = tempfile::tempdir().unwrap(); @@ -12352,6 +12434,7 @@ mod tests { .train(false) .await .unwrap(); + let readable_version = manifest_index(&dataset, "id_idx").await.index_version; hide_index_from_this_build(&mut dataset, "id_idx").await; // A fragment the index does not cover, so a bin holding it together with @@ -12367,6 +12450,8 @@ mod tests { .await .unwrap(); let before = manifest_index(&dataset, "id_idx").await; + let filter = Some("id < 3".to_string()); + assert_eq!(dataset.count_rows(filter.clone()).await.unwrap(), 6); let metrics = compact_files(&mut dataset, CompactionOptions::default(), None) .await @@ -12380,15 +12465,44 @@ mod tests { let after = manifest_index(&dataset, "id_idx").await; assert_eq!(after.uuid, before.uuid); assert_eq!(after.index_version, before.index_version); + assert_eq!(after.fields, before.fields); + assert_eq!(after.index_details, before.index_details); + // Every fragment it covered was rewritten, and it claims none of the replacements. + assert!(after.fragment_bitmap.as_ref().unwrap().is_empty()); + assert_eq!(dataset.count_rows(filter.clone()).await.unwrap(), 6); + + // A writer that can read the index rebuilds its coverage. + let mut readable = after.clone(); + readable.index_version = readable_version; + dataset + .apply_commit( + Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![readable], + removed_indices: vec![after], + }, + None, + ), + &Default::default(), + &Default::default(), + ) + .await + .unwrap(); + dataset + .create_index_builder(&["id"], IndexType::BTree, &btree_params) + .name("id_idx".to_string()) + .replace(true) + .await + .unwrap(); + let rebuilt = manifest_index(&dataset, "id_idx").await; let live = dataset .get_fragments() .iter() .map(|f| f.id() as u32) .collect::(); - assert!( - !after.fragment_bitmap.unwrap().is_disjoint(&live), - "the surviving index covers only fragments the rewrite deleted" - ); + assert_eq!(rebuilt.fragment_bitmap.unwrap(), live); + assert_eq!(dataset.count_rows(filter).await.unwrap(), 6); } /// Without stable row ids a rewrite moves every row address, so each index @@ -12398,8 +12512,9 @@ mod tests { /// longer exist. Those fragments are held back from the plan instead; the /// rest of the table still compacts. /// - /// The stable-row-id case is the test above: there the fragment-reuse index - /// repairs the coverage afterwards, so nothing has to be held back. + /// The stable-row-id case is the test above: there nothing is held back, and the + /// index instead loses coverage of the rewritten fragments until a compatible + /// writer rebuilds it. #[rstest] #[case::newer_version(UnreadableIndexKind::NewerVersion)] #[case::unknown_type(UnreadableIndexKind::UnknownType)] diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 245c861de72..1e2f6631aee 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -485,7 +485,9 @@ async fn merge_scalar_indices<'a>( // Scalar Index that expos an N:1 segment-merge primitive reachable without // rescanning the dataset let has_segment_merge_primitive = matches!(index_type, IndexType::BTree | IndexType::NGram); - let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?; + let frag_reuse_index = dataset + .frag_reuse_index_for(reference_idx, &NoOpMetricsCollector) + .await?; let ngram_requires_rebuild = index_type == IndexType::NGram && frag_reuse_index.as_ref().is_some_and(|frag_reuse_index| { fragment_reuse_affects_segments(frag_reuse_index, selected_old_indices.iter().copied()) diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index b3bd6178932..879d09a4da9 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -265,7 +265,7 @@ impl<'a> CreateIndexBuilder<'a> { let indices = load_all_indices(self.dataset).await?; let fri = self .dataset - .open_frag_reuse_index(&NoOpMetricsCollector) + .frag_reuse_index_for_row_id_entries(&NoOpMetricsCollector) .await?; // Read without consuming: a failed build must leave the requested name in // place so a retry commits under it instead of an auto-generated one. diff --git a/rust/lance/src/index/scalar.rs b/rust/lance/src/index/scalar.rs index 894080e25a8..bbdb369774c 100644 --- a/rust/lance/src/index/scalar.rs +++ b/rust/lance/src/index/scalar.rs @@ -62,7 +62,7 @@ use lance_index::scalar::{ inverted::INVERT_LIST_FILE, lance_format::LanceIndexStore, }; use lance_index::{IndexCriteria, IndexType}; -use lance_table::format::{Fragment, IndexMetadata}; +use lance_table::format::{Fragment, IndexMetadata, MAX_JSON_INDEX_NESTING}; use log::info; use prost::{Message, Name}; use tracing::instrument; @@ -307,19 +307,34 @@ impl IndexDetails { /// Returns whether this build has a reader for the complete declared type. pub(crate) fn has_reader(&self) -> bool { - let Some((_, details_type_name)) = self.0.type_url.rsplit_once('/') else { - return false; - }; - if details_type_name.is_empty() || details_type_name.starts_with('.') { - return false; + let mut details = self.0.clone(); + for _ in 0..=MAX_JSON_INDEX_NESTING { + let Some((_, details_type_name)) = details.type_url.rsplit_once('/') else { + return false; + }; + if details_type_name.is_empty() || details_type_name.starts_with('.') { + return false; + } + if details_type_name + .eq_ignore_ascii_case(&lance_index::pb::JsonIndexDetails::full_name()) + { + match lance_index::pb::JsonIndexDetails::decode(details.value.as_slice()) + .ok() + .and_then(|json| json.target_details) + { + Some(target) => details = Arc::new(target), + None => return false, + } + continue; + } + return details_type_name.eq_ignore_ascii_case(&VectorIndexDetails::full_name()) + // MemWAL flush briefly wrote this pre-`pb` package name. Keep that + // exact historical native identity readable without accepting any + // other message that merely shares the VectorIndexDetails suffix. + || details_type_name.eq_ignore_ascii_case("lance.index.VectorIndexDetails") + || SCALAR_INDEX_PLUGIN_REGISTRY.supports_details(details.as_ref()); } - - details_type_name.eq_ignore_ascii_case(&VectorIndexDetails::full_name()) - // MemWAL flush briefly wrote this pre-`pb` package name. Keep that - // exact historical native identity readable without accepting any - // other message that merely shares the VectorIndexDetails suffix. - || details_type_name.eq_ignore_ascii_case("lance.index.VectorIndexDetails") - || SCALAR_INDEX_PLUGIN_REGISTRY.supports_details(self.0.as_ref()) + false } /// Returns the index version @@ -572,7 +587,7 @@ pub async fn open_scalar_index( let index_details = fetch_index_details(dataset, column, index).await?; let plugin = SCALAR_INDEX_PLUGIN_REGISTRY.get_plugin_by_details(index_details.as_ref())?; - let frag_reuse_index = dataset.open_frag_reuse_index(metrics).await?; + let frag_reuse_index = dataset.frag_reuse_index_for(index, metrics).await?; let index_cache = dataset .index_cache @@ -613,7 +628,12 @@ pub(crate) async fn cached_scalar_index_container( dataset: &Dataset, uuid: &Uuid, ) -> Option> { - let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let index_meta = dataset.load_index(uuid).await.ok().flatten()?; + let frag_reuse_uuid = dataset + .frag_reuse_index_for(&index_meta, &NoOpMetricsCollector) + .await + .ok()? + .map(|index| index.uuid); let index_cache = dataset .index_cache .for_index(uuid, frag_reuse_uuid.as_ref()); diff --git a/rust/lance/src/index/scalar/ngram.rs b/rust/lance/src/index/scalar/ngram.rs index 11f0d5a0f52..51fbfd1e6f1 100644 --- a/rust/lance/src/index/scalar/ngram.rs +++ b/rust/lance/src/index/scalar/ngram.rs @@ -67,7 +67,9 @@ pub(in crate::index) async fn merge_segments( let segment_refs = segments.iter().collect::>(); let new_uuid = Uuid::new_v4(); - let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?; + let frag_reuse_index = dataset + .frag_reuse_index_for(&segments[0], &NoOpMetricsCollector) + .await?; let has_retired_coverage = segments.iter().any(|segment| { segment .deleted_fragment_bitmap(&dataset.fragment_bitmap) @@ -160,7 +162,14 @@ pub(in crate::index) async fn open_and_merge_segments( ) -> Result { let segments = segments.iter().map(|&s| s.clone()).collect::>(); let segment_stores = collect_ngram_segment_stores(dataset, &segments).await?; - let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?; + let frag_reuse_index = match segments.first() { + Some(segment) => { + dataset + .frag_reuse_index_for(segment, &NoOpMetricsCollector) + .await? + } + None => None, + }; let frag_reuse_index = frag_reuse_index .map(|index| Arc::new(CompactFragReuseIndexHandle(index)) as Arc); NGramIndex::merge_segments_with_remapper( diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 061a2f39d85..a644a0a0940 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -1641,7 +1641,7 @@ pub(crate) async fn open_vector_index( let mut last_stage: Option> = None; - let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let frag_reuse_uuid = frag_reuse_index.as_ref().map(|index| index.uuid); for stg in vec_idx.stages.iter().rev() { match stg.stage.as_ref() { @@ -1740,7 +1740,7 @@ pub(crate) async fn open_vector_index_v2( let index_metadata: lance_index::IndexMetadata = serde_json::from_str(index_metadata)?; let distance_type = DistanceType::try_from(index_metadata.distance_type.as_str())?; - let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let frag_reuse_uuid = frag_reuse_index.as_ref().map(|index| index.uuid); // Load the index metadata to get the correct index directory let index_meta = dataset .load_index(uuid) @@ -1938,7 +1938,7 @@ pub async fn initialize_vector_index( let new_uuid = Uuid::new_v4(); let frag_reuse_index = target_dataset - .open_frag_reuse_index(&NoOpMetricsCollector) + .frag_reuse_index_for_row_id_entries(&NoOpMetricsCollector) .await?; let summary = build_vector_index_incremental( diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index da21797e09c..3d1447927b2 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -793,7 +793,9 @@ pub(crate) async fn optimize_vector_indices_v2( let distance_type = reference_index.metric_type(); let num_partitions = ivf_model.num_partitions(); let index_type = reference_index.sub_index_type(); - let frag_reuse_index = dataset.open_frag_reuse_index(&NoOpMetricsCollector).await?; + let frag_reuse_index = dataset + .frag_reuse_index_for_row_id_entries(&NoOpMetricsCollector) + .await?; let format_version = dataset_format_version(dataset); diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 27329daf482..f055420929c 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -55,8 +55,8 @@ use crate::dataset::{ ManifestWriteConfig, NewTransactionResult, TRANSACTIONS_DIR, load_new_transactions, write_manifest_file, }; -use crate::index::DatasetIndexInternalExt; use crate::index::vector::details::infer_missing_vector_details; +use crate::index::{DatasetIndexInternalExt, stable_row_id_index_domain}; use crate::index::{index_is_usable, load_all_indices}; use crate::io::deletion::read_dataset_deletion_file; use crate::session::Session; @@ -1207,11 +1207,13 @@ pub(crate) async fn do_commit_detached_transaction( ) .await? } - _ => transaction.build_manifest( + _ => transaction.build_manifest_with_index_domain( Some(dataset.manifest.as_ref()), load_all_indices(dataset).await?.as_ref().clone(), &transaction_file, &write_config.to_build_config(), + None, + &stable_row_id_index_domain, )?, }; @@ -1559,12 +1561,13 @@ pub(crate) async fn commit_transaction( ) .await? } - _ => transaction.build_manifest_with_read_version( + _ => transaction.build_manifest_with_index_domain( Some(dataset.manifest.as_ref()), load_all_indices(&dataset).await?.as_ref().clone(), transaction_file, &write_config.to_build_config(), read_version_state, + &stable_row_id_index_domain, )?, };