fix: make FTS metadata loading retry-safe

rust/lance-core/src/error.rs

@@ -519,6 +519,17 @@ impl Error {
         NotFoundSnafu { uri: uri.into() }.build()
     }
 
+    /// Return whether this error or one of its typed sources is a missing object.
+    pub fn is_not_found(&self) -> bool {
+        match self {
+            Self::NotFound { .. } => true,
+            Self::IO { source, .. } | Self::Wrapped { error: source, .. } => {
+                error_source_is_not_found(source.as_ref())
+            }
+            _ => false,
+        }
+    }
+
     #[track_caller]
     pub fn wrapped(error: BoxedError) -> Self {
         WrappedSnafu { error }.build()
@@ -700,6 +711,17 @@ impl Error {
     }
 }
 
+fn error_source_is_not_found(source: &(dyn std::error::Error + 'static)) -> bool {
+    if let Some(error) = source.downcast_ref::<Error>() {
+        return error.is_not_found();
+    }
+    if let Some(error) = source.downcast_ref::<object_store::Error>() {
+        return matches!(error, object_store::Error::NotFound { .. })
+            || std::error::Error::source(error).is_some_and(error_source_is_not_found);
+    }
+    source.source().is_some_and(error_source_is_not_found)
+}
+
 pub trait LanceOptionExt<T> {
     /// Unwraps an option, returning an internal error if the option is None.
     ///
@@ -905,14 +927,18 @@ pub fn get_caller_location() -> &'static std::panic::Location<'static> {
 /// Wrap an error in a new error type that implements Clone
 ///
 /// This is useful when two threads/streams share a common fallible source
-/// The base error will always have the full error.  Any cloned results will
-/// only have Error::Cloned with the to_string of the base error.
+/// Definite not-found errors preserve their variant so compatibility fallbacks
+/// can distinguish them from transient I/O failures. Other cloned results use
+/// Error::Cloned with the string representation of the base error.
 pub struct CloneableError(pub Error);
 
 impl Clone for CloneableError {
     #[track_caller]
     fn clone(&self) -> Self {
-        Self(Error::cloned(self.0.to_string()))
+        match &self.0 {
+            Error::NotFound { uri, .. } => Self(Error::not_found(uri.clone())),
+            error => Self(Error::cloned(error.to_string())),
+        }
     }
 }
 
@@ -930,6 +956,19 @@ mod test {
     use super::*;
     use std::fmt;
 
+    #[test]
+    fn cloneable_error_preserves_not_found_variant() {
+        let original = CloneableError(Error::not_found("metadata.lance"));
+        let cloned = original.clone();
+        assert!(matches!(original.0, Error::NotFound { .. }));
+        assert!(matches!(cloned.0, Error::NotFound { .. }));
+
+        let original = CloneableError(Error::timeout("metadata read timed out"));
+        let cloned = original.clone();
+        assert!(matches!(original.0, Error::Timeout { .. }));
+        assert!(matches!(cloned.0, Error::Cloned { .. }));
+    }
+
     #[test]
     fn test_caller_location_capture() {
         let current_fn = get_caller_location();

rust/lance-index/src/scalar/inverted/index.rs

@@ -1218,7 +1218,7 @@ impl InvertedIndex {
                     .ok_or(Error::index("params not found in metadata".to_owned()))?;
                 Ok(serde_json::from_str::<InvertedIndexParams>(params)?)
             }
-            Err(_) => {
+            Err(error) if error.is_not_found() => {
                 // Legacy format: params live in the tokens file (see
                 // `load_legacy_index`).
                 let reader = store.open_index_file(TOKENS_FILE).await?;
@@ -1230,6 +1230,7 @@ impl InvertedIndex {
                     .transpose()?
                     .unwrap_or_default())
             }
+            Err(error) => Err(error),
         }
     }
 
@@ -7164,6 +7165,14 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn params_legacy_fallback_only_accepts_missing_metadata() {
+        assert!(Error::not_found(METADATA_FILE).is_not_found());
+        assert!(Error::wrapped(Box::new(Error::not_found(METADATA_FILE))).is_not_found());
+        assert!(!Error::timeout("metadata read timed out").is_not_found());
+        assert!(!Error::io("metadata read was denied").is_not_found());
+    }
+
     async fn write_single_partition_index(
         store: Arc<LanceIndexStore>,
         params: InvertedIndexParams,

rust/lance/src/dataset.rs

@@ -741,6 +741,7 @@ impl Dataset {
             let ds_index_cache = session.index_cache.for_dataset(uri);
             let metadata_key = crate::session::index_caches::IndexMetadataKey {
                 version: manifest_location.version,
+                store_identity: &object_store.store_prefix,
             };
             ds_index_cache
                 .insert_with_key(&metadata_key, Arc::new(indices))

rust/lance/src/dataset/fragment.rs

@@ -3748,11 +3748,11 @@ mod tests {
             // A take that hits the coverage does need the file, so it now fails with
             // a not-found error naming the missing overlay file.
             let err = frag.take(&[5], &val_only).await.unwrap_err();
-            let err = format!("{err:?}");
+            let message = format!("{err:?}");
             assert!(
-                err.contains("miss.lance") && err.to_lowercase().contains("not found"),
+                err.is_not_found() && message.contains("miss.lance"),
                 "take hitting the overlay's coverage should fail with a not-found error \
-                 for its missing file, got: {err}",
+                 for its missing file, got: {message}",
             );
         }
 
@@ -4192,11 +4192,11 @@ mod tests {
             // Projecting the overlaid `val` column does need the file, so it fails
             // with a not-found error naming the missing overlay file.
             let err = frag.take(&[0], &val_only).await.unwrap_err();
-            let err = format!("{err:?}");
+            let message = format!("{err:?}");
             assert!(
-                err.contains("valov.lance") && err.to_lowercase().contains("not found"),
+                err.is_not_found() && message.contains("valov.lance"),
                 "projecting the overlaid column should fail with a not-found error \
-                 for its missing file, got: {err}",
+                 for its missing file, got: {message}",
             );
         }
 

rust/lance/src/index.rs

@@ -1314,24 +1314,21 @@ impl DatasetIndexExt for Dataset {
     async fn load_indices(&self) -> Result<Arc<Vec<IndexMetadata>>> {
         let metadata_key = IndexMetadataKey {
             version: self.version().version,
+            store_identity: &self.object_store.store_prefix,
         };
-        let mut indices = match self.index_cache.get_with_key(&metadata_key).await {
-            Some(indices) => indices,
-            None => {
+        let mut indices = self
+            .index_cache
+            .get_or_insert_with_key(metadata_key, || async {
                 let mut loaded_indices = read_manifest_indexes(
                     &self.object_store,
                     &self.manifest_location,
                     &self.manifest,
                 )
                 .await?;
                 retain_supported_indices(&mut loaded_indices);
-                let loaded_indices = Arc::new(loaded_indices);
-                self.index_cache
-                    .insert_with_key(&metadata_key, loaded_indices.clone())
-                    .await;
-                loaded_indices
-            }
-        };
+                Ok(loaded_indices)
+            })
+            .await?;
 
         // Infer details for legacy vector indices (once per index name, concurrently).
         // This may run on indices that were opportunistically cached during Dataset::open
@@ -4350,6 +4347,48 @@ mod tests {
         assert_eq!(indices2.len(), 1);
     }
 
+    #[tokio::test]
+    async fn test_load_indices_singleflights_concurrent_cache_misses() {
+        let session = Arc::new(Session::default());
+        let write_params = WriteParams {
+            session: Some(session.clone()),
+            ..Default::default()
+        };
+        let test_dir = TempStrDir::default();
+        let schema = Arc::new(Schema::new(vec![Field::new("tag", DataType::Utf8, false)]));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(StringArray::from(vec!["a", "b", "c"]))],
+        )
+        .unwrap();
+        let mut dataset = Dataset::write(
+            RecordBatchIterator::new(vec![Ok(batch)], schema),
+            &test_dir,
+            Some(write_params),
+        )
+        .await
+        .unwrap();
+        dataset
+            .create_index(
+                &["tag"],
+                IndexType::Bitmap,
+                None,
+                &ScalarIndexParams::default(),
+                false,
+            )
+            .await
+            .unwrap();
+
+        session.index_cache.clear().await;
+        let before = session.index_cache_stats().await;
+        let results = futures::future::join_all((0..32).map(|_| dataset.load_indices())).await;
+        assert!(results.iter().all(|result| result.is_ok()));
+        let after = session.index_cache_stats().await;
+
+        assert_eq!(after.misses - before.misses, 1);
+        assert_eq!(after.hits - before.hits, 31);
+    }
+
     #[tokio::test]
     async fn test_remap_empty() {
         let data = gen_batch()
@@ -4736,6 +4775,7 @@ mod tests {
         // We commit by doing a delete("false") after replacing the cached indices.
         let metadata_key = crate::session::index_caches::IndexMetadataKey {
             version: dataset.version().version,
+            store_identity: &dataset.object_store.store_prefix,
         };
         dataset
             .index_cache

rust/lance/src/io/commit.rs

@@ -1073,6 +1073,7 @@ pub(crate) async fn commit_transaction(
                 if !indices.is_empty() {
                     let key = IndexMetadataKey {
                         version: target_version,
+                        store_identity: &dataset.object_store.store_prefix,
                     };
                     dataset
                         .index_cache

rust/lance/src/session/index_caches.rs

@@ -94,16 +94,22 @@ impl CacheKey for FragReuseIndexKey<'_> {
     }
 }
 
-#[derive(Debug)]
-pub struct IndexMetadataKey {
+#[derive(Clone, Copy, Debug)]
+pub struct IndexMetadataKey<'a> {
     pub version: u64,
+    pub store_identity: &'a str,
 }
 
-impl CacheKey for IndexMetadataKey {
+impl CacheKey for IndexMetadataKey<'_> {
     type ValueType = Vec<IndexMetadata>;
 
     fn key(&self) -> Cow<'_, str> {
-        Cow::Owned(self.version.to_string())
+        Cow::Owned(format!(
+            "{}:{}/{}",
+            self.store_identity.len(),
+            self.store_identity,
+            self.version
+        ))
     }
 
     fn type_name() -> &'static str {
@@ -145,3 +151,22 @@ impl CacheKey for ScalarIndexDetailsKey<'_> {
         "ScalarIndexDetails"
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn index_metadata_key_isolates_object_store_identity() {
+        let first = IndexMetadataKey {
+            version: 7,
+            store_identity: "s3$first-options",
+        };
+        let second = IndexMetadataKey {
+            version: 7,
+            store_identity: "s3$second-options",
+        };
+
+        assert_ne!(first.key(), second.key());
+    }
+}