From 5e59b6029088bf6a330e0b592dcb1124f8382ce8 Mon Sep 17 00:00:00 2001 From: Brendan Clement Date: Wed, 9 Sep 2026 16:22:01 -0700 Subject: [PATCH 1/2] feat: apply the fragment reuse index by index identifier domain On a stable-row-id dataset most indices store stable row ids, which a fragment reuse index must not remap: stable ids start at 0 and collide with fragment 0's row addresses. Index loading, merging and building now hand the fragment reuse index only to indices that store row addresses. Cleanup and remap_column_index treat stable-id indices as needing only their fragment coverage repaired, and the manifest build derives the FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS compatibility flag, which this build now supports. A JSON index answers in its target's domain; both domain predicates now follow the target, and malformed JSON details are an error rather than a guess. Under stable row ids the rewrite branch no longer recalculates the fragment reuse index's own coverage, which the eager path never touched. --- rust/lance-index/src/scalar/json.rs | 44 ++- rust/lance-table/src/feature_flags.rs | 77 +++-- rust/lance-table/src/format.rs | 5 +- rust/lance-table/src/format/index.rs | 297 ++++++++++++++++-- .../src/system_index/frag_reuse.rs | 27 ++ .../src/transaction/index_maintenance.rs | 5 +- .../src/transaction/manifest_build.rs | 246 ++++++++++++++- rust/lance/src/dataset/index/frag_reuse.rs | 120 ++++++- rust/lance/src/dataset/optimize/remapping.rs | 49 ++- rust/lance/src/index.rs | 149 ++++++++- rust/lance/src/index/append.rs | 4 +- rust/lance/src/index/create.rs | 2 +- rust/lance/src/index/scalar.rs | 50 ++- rust/lance/src/index/scalar/ngram.rs | 13 +- rust/lance/src/index/vector.rs | 6 +- rust/lance/src/index/vector/ivf.rs | 4 +- rust/lance/src/io/commit.rs | 9 +- 17 files changed, 957 insertions(+), 150 deletions(-) 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..423603e392c 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -56,9 +56,9 @@ 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. +/// Set by `build_manifest` whenever both hold, and lifted when either stops +/// holding. Supported since index loading honors the identifier domain; 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 +91,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 +157,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 +198,23 @@ 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 index loading honors the identifier domain of a +/// fragment reuse index on a stable-row-id table. 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 the domain-aware loading 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,26 @@ 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 domain-aware index loading must refuse the table for + /// reading and writing, or it would remap stable row ids as if they were + /// row addresses; 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 +354,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 +391,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/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..6139607fd70 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. A build that applies the fragment reuse index to + // every index would remap stable row ids as if they were row addresses; + // reading returns wrong rows and optimizing indices persists them. + 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..cc968468bbf 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,11 @@ 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. It is + /// applied only to indices that store row addresses, never to entries that + /// are stable row ids, and such a table sets a reader feature flag that + /// older Lance versions refuse. /// * Says nothing about deletion files, source-value changes, or whether an /// address belongs to this table or branch. /// @@ -84,6 +87,10 @@ 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, 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/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 491ad893063..b8fc2393b96 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,17 @@ 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. RemapResult::Keep(new_id) => IndexMetadata { uuid: new_id, name: curr_index_meta.name.clone(), 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, )?, }; From 670460d92ad8bd6650519e5a72bc8376a1db9c70 Mon Sep 17 00:00:00 2001 From: Brendan Clement Date: Wed, 9 Sep 2026 22:03:14 -0700 Subject: [PATCH 2/2] feat: enable the fragment reuse index on stable row id datasets Compaction with defer_index_remap no longer rejects stable row ids. The row addresses the fragment reuse index records are derived from the fragment metadata (every physical row minus the deletion vector), which is what the scan yields without stable row ids, while the row id sequences are rechunked as before. Closes #8977 --- python/python/tests/test_optimize.py | 29 +- rust/lance-table/src/feature_flags.rs | 19 +- rust/lance-table/src/format/manifest.rs | 17 +- .../src/transaction/manifest_build.rs | 6 +- rust/lance/src/dataset/index/frag_reuse.rs | 16 +- rust/lance/src/dataset/optimize.rs | 1481 +++++++++++++++-- rust/lance/src/dataset/optimize/remapping.rs | 5 +- 7 files changed, 1393 insertions(+), 180 deletions(-) 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-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 423603e392c..9680f3f76da 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -56,8 +56,9 @@ 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. /// -/// Set by `build_manifest` whenever both hold, and lifted when either stops -/// holding. Supported since index loading honors the identifier domain; see +/// 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 @@ -198,11 +199,10 @@ fn mark_supported(flags: &mut u64, flag: u64, feature_enabled: bool) { } /// The feature-flag bits this build understands, given whether overlay support -/// is enabled and whether index loading honors the identifier domain of a -/// fragment reuse index on a stable-row-id table. 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 the domain-aware loading does with such a table. +/// 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( @@ -320,9 +320,8 @@ mod tests { use super::*; use crate::format::BasePath; - /// A build without domain-aware index loading must refuse the table for - /// reading and writing, or it would remap stable row ids as if they were - /// row addresses; this build accepts it. + /// 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_gating() { use crate::format::{DataStorageFormat, Manifest}; 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/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 6139607fd70..701af9ac03f 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1496,9 +1496,9 @@ impl Transaction { manifest.reader_feature_flags |= FLAG_COVERED_INDEX_METADATA; manifest.writer_feature_flags |= FLAG_COVERED_INDEX_METADATA; } - // Derived the same way. A build that applies the fragment reuse index to - // every index would remap stable row ids as if they were row addresses; - // reading returns wrong rows and optimizing indices persists them. + // 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() diff --git a/rust/lance/src/dataset/index/frag_reuse.rs b/rust/lance/src/dataset/index/frag_reuse.rs index cc968468bbf..2a234b56486 100644 --- a/rust/lance/src/dataset/index/frag_reuse.rs +++ b/rust/lance/src/dataset/index/frag_reuse.rs @@ -44,10 +44,9 @@ impl Dataset { /// stale, for example after every row of the destination fragment is deleted. /// * Not every compaction records an FRI: it requires `defer_index_remap`, /// and fresh index-free tables do not receive one automatically. - /// * With stable row ids the FRI still describes physical addresses. It is - /// applied only to indices that store row addresses, never to entries that - /// are stable row ids, and such a table sets a reader feature flag that - /// older Lance versions refuse. + /// * 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. /// @@ -87,10 +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, 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. +/// 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, 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 b8fc2393b96..32d3edfda07 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -357,7 +357,8 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { 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. + // 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(), @@ -368,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 {