diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 21aadf7d475..07795209331 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1206,6 +1206,33 @@ impl InvertedIndex { self.partitions.len() == 1 && self.partitions[0].is_legacy() } + /// Read only the index's [`InvertedIndexParams`], + /// Contains more complete info than manifest's lossy `InvertedIndexDetails`. + pub async fn load_params(store: &dyn IndexStore) -> Result { + match store.open_index_file(METADATA_FILE).await { + Ok(reader) => { + let params = reader + .schema() + .metadata + .get("params") + .ok_or(Error::index("params not found in metadata".to_owned()))?; + Ok(serde_json::from_str::(params)?) + } + Err(_) => { + // Legacy format: params live in the tokens file (see + // `load_legacy_index`). + let reader = store.open_index_file(TOKENS_FILE).await?; + Ok(reader + .schema() + .metadata + .get("tokenizer") + .map(|s| serde_json::from_str::(s)) + .transpose()? + .unwrap_or_default()) + } + } + } + pub async fn load( store: Arc, frag_reuse_index: Option>, diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 86682004f04..b12a7198e5a 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -4113,3 +4113,50 @@ async fn test_manifest_read_recovers_from_stale_size() { assert_eq!(indices.len(), 1); assert_eq!(indices[0].name, "id_idx"); } + +/// `load_segment_params` must match the fully opened segment's params, +/// including `custom_stop_words` — the field `InvertedIndexDetails` loses. +#[tokio::test] +async fn test_load_segment_params_full_fidelity() { + use crate::index::DatasetIndexInternalExt; + use lance_index::metrics::NoOpMetricsCollector; + use lance_index::scalar::inverted::InvertedIndex; + + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![Field::new("text", DataType::Utf8, false)]).into(), + vec![Arc::new(StringArray::from(vec![ + "the quick brown fox", + "lazy dogs sleep", + ]))], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(stream, "memory://test/segment_params", None) + .await + .unwrap(); + + let params = InvertedIndexParams::default().custom_stop_words(Some(vec!["quick".to_string()])); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let segments = crate::index::scalar::load_segments(&dataset, "text") + .await + .unwrap() + .expect("FTS index segments"); + let read = crate::index::scalar::load_segment_params(&dataset, &segments[0]) + .await + .unwrap(); + + let generic = dataset + .open_generic_index("text", &segments[0].uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let opened = generic + .as_any() + .downcast_ref::() + .expect("inverted index"); + assert_eq!(&read, opened.params()); +} diff --git a/rust/lance/src/index/scalar.rs b/rust/lance/src/index/scalar.rs index d31b96c9202..2f346163974 100644 --- a/rust/lance/src/index/scalar.rs +++ b/rust/lance/src/index/scalar.rs @@ -11,7 +11,7 @@ pub(crate) mod inverted; pub(crate) mod label_list; pub(crate) mod zonemap; -pub use inverted::{load_segment_details, load_segments}; +pub use inverted::{load_segment_details, load_segment_params, load_segments}; pub use crate::index::scalar_logical::{LogicalScalarIndex, load_named_scalar_segments}; diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index b41dfa562b9..426d104c912 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -12,7 +12,7 @@ use lance_core::ROW_ID; use lance_index::metrics::NoOpMetricsCollector; use lance_index::pbold::InvertedIndexDetails; use lance_index::scalar::index_files_to_table; -use lance_index::scalar::inverted::InvertedIndex; +use lance_index::scalar::inverted::{InvertedIndex, InvertedIndexParams}; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; use lance_table::format::IndexMetadata; @@ -220,6 +220,15 @@ pub async fn load_segment_details( }) } +/// Read one segment's [`InvertedIndexParams`] +pub async fn load_segment_params( + dataset: &Dataset, + segment: &IndexMetadata, +) -> Result { + let store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?; + InvertedIndex::load_params(&store).await +} + #[cfg(test)] mod tests { use super::*;