diff --git a/rust/lance-index/src/scalar/fmindex.rs b/rust/lance-index/src/scalar/fmindex.rs index aa07e550b52..b55579fda7d 100644 --- a/rust/lance-index/src/scalar/fmindex.rs +++ b/rust/lance-index/src/scalar/fmindex.rs @@ -1642,6 +1642,11 @@ impl ScalarIndex for FMIndexScalarIndex { )), } } + + fn results_are_row_addresses(&self) -> bool { + true + } + fn can_remap(&self) -> bool { false } diff --git a/rust/lance-table/src/format/index.rs b/rust/lance-table/src/format/index.rs index 770cd983a86..6f9ed005a8d 100644 --- a/rust/lance-table/src/format/index.rs +++ b/rust/lance-table/src/format/index.rs @@ -143,6 +143,7 @@ impl IndexMetadata { self.index_details.as_ref().is_some_and(|details| { details.type_url.ends_with("ZoneMapIndexDetails") || details.type_url.ends_with("BloomFilterIndexDetails") + || details.type_url.ends_with("FMIndexDetails") }) } @@ -596,6 +597,21 @@ mod tests { } } + #[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::btree("type.googleapis.com/lance.table.BTreeIndexDetails", false)] + fn test_results_are_row_addrs(#[case] type_url: &str, #[case] expected: bool) { + 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: Vec::new(), + })); + + assert_eq!(metadata.results_are_row_addrs(), expected); + } + #[rstest] #[case::empty_is_valid(vec![7], vec![], None)] // mem_wal and frag_reuse both commit no fields at all; a bare diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 5123f389625..1b61af5ae22 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -4020,38 +4020,52 @@ mod tests { /// Regression test for https://github.com/lance-format/lance/issues/8076 /// - /// A zone map or bloom filter index reports matches as physical row addresses, so + /// Zone map, bloom filter, and FM indices report matches as physical row addresses, so /// compaction invalidates it even under stable row ids. Reusing it for the rewritten /// fragments made a filtered scan fail with an internal error (a fragment referenced /// by the index no longer existed) or, once translation tolerated that, silently drop /// every match. #[rstest] - #[case::zone_map(BuiltinIndexType::ZoneMap, IndexType::ZoneMap)] - #[case::bloom_filter(BuiltinIndexType::BloomFilter, IndexType::BloomFilter)] + #[case::zone_map(BuiltinIndexType::ZoneMap, IndexType::ZoneMap, "i", "i > 0", 199)] + #[case::bloom_filter(BuiltinIndexType::BloomFilter, IndexType::BloomFilter, "i", "i = 0", 1)] + #[case::fm( + BuiltinIndexType::Fm, + IndexType::Fm, + "text", + "contains(text, 'needle')", + 100 + )] #[tokio::test] async fn test_addr_domain_index_after_compaction_with_stable_row_ids( #[case] builtin: BuiltinIndexType, #[case] index_type: IndexType, + #[case] indexed_column: &str, + #[case] query: &str, + #[case] expected_rows: usize, ) { - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("i".to_owned()))); - let mut dataset = Dataset::write( - data_gen.batch(200), - "memory://test/table", - Some(WriteParams { - enable_stable_row_ids: true, - max_rows_per_file: 100, // 2 fragments, so compaction has something to merge - ..Default::default() - }), - ) - .await - .unwrap(); + let mut dataset = lance_datagen::gen_batch() + .col("i", lance_datagen::array::step::()) + .col( + "text", + lance_datagen::array::cycle_utf8_literals(&["needle", "haystack"]), + ) + .into_ram_dataset_with_params( + FragmentCount::from(2), + FragmentRowCount::from(100), + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 100, + ..Default::default() + }), + ) + .await + .unwrap(); dataset .create_index( - &["i"], + &[indexed_column], index_type, - None, + Some("addr_idx".to_string()), &ScalarIndexParams::for_builtin(builtin), false, ) @@ -4071,7 +4085,7 @@ mod tests { .await .unwrap() .iter() - .find(|index| index.fields == vec![0]) + .find(|index| index.name == "addr_idx") .expect("index must survive compaction") .clone(); assert!( @@ -4084,9 +4098,9 @@ mod tests { // Every fragment therefore falls back to a full scan, and the filter is answered // in full. let mut scanner = dataset.scan(); - scanner.filter("i > 0").unwrap(); + scanner.filter(query).unwrap(); let matched = scanner.try_into_batch().await.unwrap(); - assert_eq!(matched.num_rows(), 199); + assert_eq!(matched.num_rows(), expected_rows); } // Regression test for https://github.com/lancedb/lance/issues/6161 diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8a4fb75a399..8e87c2ffc9c 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -11663,6 +11663,82 @@ mod test { .unwrap(); } + #[tokio::test] + async fn test_fm_index_with_stable_row_ids() { + let batch = arrow_array::record_batch!( + ( + "text", + Utf8, + [ + "alpha", + "needle in first", + "beta", + "first needle suffix", + "delta", + "needle in second", + "epsilon", + "second needle suffix" + ] + ), + ("id", Int32, [0, 1, 2, 3, 4, 5, 6, 7]) + ) + .unwrap(); + let schema = batch.schema(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let write_params = WriteParams { + max_rows_per_file: 4, + enable_stable_row_ids: true, + ..Default::default() + }; + let mut dataset = Dataset::write( + reader, + "memory://test_fm_index_with_stable_row_ids", + Some(write_params), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 2); + + let params = ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::Fm); + dataset + .create_index(&["text"], IndexType::Fm, None, ¶ms, true) + .await + .unwrap(); + + let mut indexed_scan = dataset.scan(); + indexed_scan.filter("contains(text, 'needle')").unwrap(); + let indexed_plan = indexed_scan.explain_plan(false).await.unwrap(); + assert!( + indexed_plan.contains("ScalarIndexQuery") && indexed_plan.contains("Fm"), + "expected the FM index in the plan, got:\n{indexed_plan}" + ); + let indexed = indexed_scan.try_into_batch().await.unwrap(); + + let unindexed = dataset + .scan() + .use_scalar_index(false) + .filter("contains(text, 'needle')") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let indexed_ids = indexed["id"] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + let unindexed_ids = unindexed["id"] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + + assert_eq!(unindexed_ids, BTreeSet::from([1, 3, 5, 7])); + assert_eq!(indexed_ids, unindexed_ids); + } + #[tokio::test] async fn test_ngram_regex_index_scan() { use arrow::array::AsArray; diff --git a/rust/lance/src/dataset/write/update.rs b/rust/lance/src/dataset/write/update.rs index e72002dcff3..00e57c8fad6 100644 --- a/rust/lance/src/dataset/write/update.rs +++ b/rust/lance/src/dataset/write/update.rs @@ -1444,17 +1444,30 @@ mod tests { } #[rstest] - #[case::zone_map(BuiltinIndexType::ZoneMap, "i < 100", 100)] - #[case::bloom_filter(BuiltinIndexType::BloomFilter, "i = 0", 1)] + #[case::zone_map(BuiltinIndexType::ZoneMap, IndexType::ZoneMap, "i", "i < 100", 100)] + #[case::bloom_filter(BuiltinIndexType::BloomFilter, IndexType::BloomFilter, "i", "i = 0", 1)] + #[case::fm( + BuiltinIndexType::Fm, + IndexType::Fm, + "text", + "contains(text, 'needle')", + 50 + )] #[tokio::test] async fn test_addr_domain_index_does_not_cover_rewritten_update_fragment( - #[case] index_type: BuiltinIndexType, + #[case] builtin: BuiltinIndexType, + #[case] index_type: IndexType, + #[case] indexed_column: &str, #[case] query: &str, #[case] expected_rows: usize, ) { let mut dataset = lance_datagen::gen_batch() .col("i", lance_datagen::array::step::()) .col("category", lance_datagen::array::step::()) + .col( + "text", + lance_datagen::array::cycle_utf8_literals(&["needle", "haystack"]), + ) .into_ram_dataset_with_params( FragmentCount::from(1), FragmentRowCount::from(100), @@ -1469,10 +1482,10 @@ mod tests { dataset .create_index( - &["i"], - IndexType::Scalar, - Some("i_idx".to_string()), - &ScalarIndexParams::for_builtin(index_type), + &[indexed_column], + index_type, + Some("addr_idx".to_string()), + &ScalarIndexParams::for_builtin(builtin), true, ) .await @@ -1500,7 +1513,10 @@ mod tests { .new_dataset; let indices = dataset.load_indices().await.unwrap(); - let index = indices.iter().find(|index| index.name == "i_idx").unwrap(); + let index = indices + .iter() + .find(|index| index.name == "addr_idx") + .unwrap(); assert_eq!( index .fragment_bitmap