feat: auto-migrate old index metadata (#3428)

Follow up to #3377. That PR made
`index_statistics()` error by default. This ended up being a footgun for
some users who rely heavily on that method. So instead of forcing the
user to do the migration themself, we do it for them. It can be disabled
using an environment variable.

rust/lance-core/src/utils.rs

@@ -8,6 +8,7 @@ pub mod deletion;
 pub mod futures;
 pub mod hash;
 pub mod mask;
+pub mod parse;
 pub mod path;
 pub mod testing;
 pub mod tokio;

rust/lance-core/src/utils/parse.rs

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+/// Parse a string into a boolean value.
+pub fn str_is_truthy(val: &str) -> bool {
+    val.eq_ignore_ascii_case("1")
+        | val.eq_ignore_ascii_case("true")
+        | val.eq_ignore_ascii_case("on")
+        | val.eq_ignore_ascii_case("yes")
+        | val.eq_ignore_ascii_case("y")
+}

rust/lance-io/src/object_store.rs

@@ -17,6 +17,7 @@ use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use deepsize::DeepSizeOf;
 use futures::{future, stream::BoxStream, StreamExt, TryStreamExt};
+use lance_core::utils::parse::str_is_truthy;
 use lance_core::utils::tokio::get_num_compute_intensive_cpus;
 use object_store::aws::{
     AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential, AwsCredentialProvider,
@@ -1039,14 +1040,6 @@ fn infer_block_size(scheme: &str) -> usize {
     }
 }
 
-fn str_is_truthy(val: &str) -> bool {
-    val.eq_ignore_ascii_case("1")
-        | val.eq_ignore_ascii_case("true")
-        | val.eq_ignore_ascii_case("on")
-        | val.eq_ignore_ascii_case("yes")
-        | val.eq_ignore_ascii_case("y")
-}
-
 /// Attempt to create a Url from given table location.
 ///
 /// The location could be:

rust/lance/src/dataset.rs

@@ -4395,10 +4395,13 @@ mod tests {
         let validate_res = dataset.validate().await;
         assert!(validate_res.is_err());
         assert_eq!(dataset.load_indices().await.unwrap()[0].name, "vector_idx");
-        assert!(dataset.index_statistics("vector_idx").await.is_err());
 
-        // Force a migration
-        dataset.delete("false").await.unwrap();
+        // Calling index statistics will force a migration
+        let stats = dataset.index_statistics("vector_idx").await.unwrap();
+        let stats: serde_json::Value = serde_json::from_str(&stats).unwrap();
+        assert_eq!(stats["num_indexed_fragments"], 2);
+
+        dataset.checkout_latest().await.unwrap();
         dataset.validate().await.unwrap();
 
         let indices = dataset.load_indices().await.unwrap();
@@ -4408,10 +4411,6 @@ mod tests {
         }
         assert_eq!(get_bitmap(&indices[0]), vec![0]);
         assert_eq!(get_bitmap(&indices[1]), vec![1]);
-
-        let stats = dataset.index_statistics("vector_idx").await.unwrap();
-        let stats: serde_json::Value = serde_json::from_str(&stats).unwrap();
-        assert_eq!(stats["num_indexed_fragments"], 2);
     }
 
     #[rstest]

rust/lance/src/index.rs

@@ -5,12 +5,13 @@
 //!
 
 use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 use arrow_schema::DataType;
 use async_trait::async_trait;
 use futures::{stream, StreamExt, TryStreamExt};
 use itertools::Itertools;
+use lance_core::utils::parse::str_is_truthy;
 use lance_file::reader::FileReader;
 use lance_file::v2;
 use lance_file::v2::reader::FileReaderOptions;
@@ -68,6 +69,17 @@ use self::append::merge_indices;
 use self::scalar::build_scalar_index;
 use self::vector::{build_vector_index, VectorIndexParams, LANCE_VECTOR_INDEX};
 
+// Whether to auto-migrate a dataset when we encounter corruption.
+fn auto_migrate_corruption() -> bool {
+    static LANCE_AUTO_MIGRATION: OnceLock<bool> = OnceLock::new();
+    *LANCE_AUTO_MIGRATION.get_or_init(|| {
+        std::env::var("LANCE_AUTO_MIGRATION")
+            .ok()
+            .map(|s| str_is_truthy(&s))
+            .unwrap_or(true)
+    })
+}
+
 /// Builds index.
 #[async_trait]
 pub trait IndexBuilder {
@@ -590,7 +602,8 @@ impl DatasetIndexExt for Dataset {
         let index_type = indices[0].index_type().to_string();
 
         let indexed_fragments_per_delta = self.indexed_fragments(index_name).await?;
-        let num_indexed_rows_per_delta = indexed_fragments_per_delta
+
+        let res = indexed_fragments_per_delta
             .iter()
             .map(|frags| {
                 let mut sum = 0;
@@ -604,18 +617,49 @@ impl DatasetIndexExt for Dataset {
                 }
                 Ok(sum)
             })
-            .collect::<Result<Vec<_>>>()?;
+            .collect::<Result<Vec<_>>>();
+
+        async fn migrate_and_recompute(ds: &Dataset, index_name: &str) -> Result<String> {
+            let mut ds = ds.clone();
+            log::warn!(
+                "Detecting out-dated fragment metadata, migrating dataset. \
+                        To disable migration, set LANCE_AUTO_MIGRATION=false"
+            );
+            ds.delete("false").await.map_err(|err| {
+                Error::Execution {
+                    message: format!("Failed to migrate dataset while calculating index statistics. \
+                            To disable migration, set LANCE_AUTO_MIGRATION=false. Original error: {}", err),
+                    location: location!(),
+                }
+            })?;
+            ds.index_statistics(index_name).await
+        }
+
+        let num_indexed_rows_per_delta = match res {
+            Ok(rows) => rows,
+            Err(Error::Internal { message, .. })
+                if auto_migrate_corruption() && message.contains("trigger a single write") =>
+            {
+                return migrate_and_recompute(self, index_name).await;
+            }
+            Err(e) => return Err(e),
+        };
 
         let mut fragment_ids = HashSet::new();
         for frags in indexed_fragments_per_delta.iter() {
             for frag in frags.iter() {
                 if !fragment_ids.insert(frag.id) {
-                    return Err(Error::Internal {
-                        message: "Overlap in indexed fragments. Please upgrade to lance >= 0.23.0 \
+                    if auto_migrate_corruption() {
+                        return migrate_and_recompute(self, index_name).await;
+                    } else {
+                        return Err(Error::Internal {
+                            message:
+                                "Overlap in indexed fragments. Please upgrade to lance >= 0.23.0 \
                                   and trigger a single write to fix this"
-                            .to_string(),
-                        location: location!(),
-                    });
+                                    .to_string(),
+                            location: location!(),
+                        });
+                    }
                 }
             }
         }