fix: add versioning and bypass broken row counts (lancedb#1534)

Adds a new feature: `WriterVersion` in the manifest.

Also fixes two bugs:
* lancedb#1531 Adds bypass logic to read correct `physical_rows` when old
version is detected.
  Will update the value on write as needed.
* When `migrate_fragments` runs, we do it on the new fragments and not
the old,
  fixing a data loss issue.

Fixes lancedb#1531
Fixes lancedb#1535

Author: wjones127

.gitignore

@@ -51,7 +51,6 @@ python/lance/_*.cpp
 
 *.parquet
 *.parq
-*.lance
 
 python/thirdparty/arrow/
 python/wheels
@@ -72,7 +71,9 @@ integration/duckdb/lance.duckdb_extension.*.zip
 
 notebooks/lance.duckdb_extension
 notebooks/sift
+notebooks/image_data/data
 benchmarks/sift/sift
+benchmarks/sift/sift.lance
 benchmarks/sift/lance_ivf*.csv
 **/sift.tar.gz
 
@@ -87,9 +88,9 @@ wheelhouse
 # Rust
 target
 Cargo.lock
-data
 
 # c++ lsp
 .ccls-cache/
 
-python/venv
+python/venv
+test_data/venv

protos/format.proto

@@ -64,6 +64,22 @@ message Manifest {
   // Schema metadata.
   map<string, bytes> metadata = 5;
 
+  message WriterVersion {
+    // The name of the library that created this file.
+    string library = 1;
+    // The version of the library that created this file. Because we cannot assume
+    // that the library is semantically versioned, this is a string. However, if it
+    // is semantically versioned, it should be a valid semver string without any 'v'
+    // prefix. For example: `2.0.0`, `2.0.0-rc.1`.
+    string version = 2;
+  }
+
+  // The version of the writer that created this file.
+  //
+  // This information may be used to detect whether the file may have known bugs
+  // associated with that writer.
+  WriterVersion writer_version = 13;
+
   // If presented, the file position of the index metadata.
   optional uint64 index_section = 6;
 

rust/lance-core/src/format.rs

@@ -25,7 +25,7 @@ mod page_table;
 
 pub use fragment::*;
 pub use index::Index;
-pub use manifest::Manifest;
+pub use manifest::{Manifest, WriterVersion};
 pub use metadata::{Metadata, StatisticsMetadata};
 pub use page_table::{PageInfo, PageTable};
 

rust/lance-core/src/format/manifest.rs

@@ -40,6 +40,9 @@ pub struct Manifest {
     /// Dataset version
     pub version: u64,
 
+    /// Version of the writer library that wrote this manifest.
+    pub writer_version: Option<WriterVersion>,
+
     /// Fragments, the pieces to build the dataset.
     pub fragments: Arc<Vec<Fragment>>,
 
@@ -73,6 +76,7 @@ impl Manifest {
         Self {
             schema: schema.clone(),
             version: 1,
+            writer_version: Some(WriterVersion::default()),
             fragments,
             version_aux_data: 0,
             index_section: None,
@@ -93,6 +97,7 @@ impl Manifest {
         Self {
             schema: schema.clone(),
             version: previous.version + 1,
+            writer_version: Some(WriterVersion::default()),
             fragments,
             version_aux_data: 0,
             index_section: None, // Caller should update index if they want to keep them.
@@ -171,6 +176,34 @@ impl Manifest {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct WriterVersion {
+    pub library: String,
+    pub version: String,
+}
+
+impl WriterVersion {
+    /// Try to parse the version string as a semver string. Returns None if
+    /// not successful.
+    pub fn semver(&self) -> Option<(u32, u32, u32, Option<&str>)> {
+        let mut parts = self.version.split('.');
+        let major = parts.next().unwrap_or("0").parse().ok()?;
+        let minor = parts.next().unwrap_or("0").parse().ok()?;
+        let patch = parts.next().unwrap_or("0").parse().ok()?;
+        let tag = parts.next();
+        Some((major, minor, patch, tag))
+    }
+}
+
+impl Default for WriterVersion {
+    fn default() -> Self {
+        Self {
+            library: "lance".to_string(),
+            version: env!("CARGO_PKG_VERSION").to_string(),
+        }
+    }
+}
+
 impl ProtoStruct for Manifest {
     type Proto = pb::Manifest;
 }
@@ -182,9 +215,17 @@ impl From<pb::Manifest> for Manifest {
             let nanos = ts.nanos as u128;
             sec + nanos
         });
+        // We only use the writer version if it is fully set.
+        let writer_version = match p.writer_version {
+            Some(pb::manifest::WriterVersion { library, version }) => {
+                Some(WriterVersion { library, version })
+            }
+            _ => None,
+        };
         Self {
             schema: Schema::from((&p.fields, p.metadata)),
             version: p.version,
+            writer_version,
             fragments: Arc::new(p.fragments.iter().map(Fragment::from).collect()),
             version_aux_data: p.version_aux_data as usize,
             index_section: p.index_section.map(|i| i as usize),
@@ -218,6 +259,13 @@ impl From<&Manifest> for pb::Manifest {
         Self {
             fields,
             version: m.version,
+            writer_version: m
+                .writer_version
+                .as_ref()
+                .map(|wv| pb::manifest::WriterVersion {
+                    library: wv.library.clone(),
+                    version: wv.version.clone(),
+                }),
             fragments: m.fragments.iter().map(pb::DataFragment::from).collect(),
             metadata,
             version_aux_data: m.version_aux_data as u64,
@@ -231,3 +279,24 @@ impl From<&Manifest> for pb::Manifest {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_writer_version() {
+        let wv = WriterVersion::default();
+        assert_eq!(wv.library, "lance");
+        assert_eq!(wv.version, env!("CARGO_PKG_VERSION"));
+        assert_eq!(
+            wv.semver(),
+            Some((
+                env!("CARGO_PKG_VERSION_MAJOR").parse().unwrap(),
+                env!("CARGO_PKG_VERSION_MINOR").parse().unwrap(),
+                env!("CARGO_PKG_VERSION_PATCH").parse().unwrap(),
+                None
+            ))
+        );
+    }
+}

rust/lance/src/dataset.rs

@@ -1507,10 +1507,11 @@ mod tests {
     use arrow_schema::{DataType, Field, Fields as ArrowFields, Schema as ArrowSchema};
     use arrow_select::take::take;
     use futures::stream::TryStreamExt;
+    use lance_core::format::WriterVersion;
     use lance_index::vector::DIST_COL;
     use lance_linalg::distance::MetricType;
     use lance_testing::datagen::generate_random_array;
-    use tempfile::tempdir;
+    use tempfile::{tempdir, TempDir};
 
     // Used to validate that futures returned are Send.
     fn require_send<T: Send>(t: T) -> T {
@@ -1561,6 +1562,10 @@ mod tests {
 
         let actual_ds = Dataset::open(test_uri).await.unwrap();
         assert_eq!(actual_ds.version().version, 1);
+        assert_eq!(
+            actual_ds.manifest.writer_version,
+            Some(WriterVersion::default())
+        );
         let actual_schema = ArrowSchema::from(actual_ds.schema());
         assert_eq!(&actual_schema, schema.as_ref());
 
@@ -3188,4 +3193,186 @@ mod tests {
             &[1],
         );
     }
+
+    fn copy_dir_all(
+        src: impl AsRef<std::path::Path>,
+        dst: impl AsRef<std::path::Path>,
+    ) -> std::io::Result<()> {
+        use std::fs;
+        fs::create_dir_all(&dst)?;
+        for entry in fs::read_dir(src)? {
+            let entry = entry?;
+            let ty = entry.file_type()?;
+            if ty.is_dir() {
+                copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name()))?;
+            } else {
+                fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?;
+            }
+        }
+        Ok(())
+    }
+
+    /// Copies a test dataset into a temporary directory, returning the tmpdir.
+    ///
+    /// The `table_path` should be relative to `test_data/` at the root of the
+    /// repo.
+    fn copy_test_data_to_tmp(table_path: &str) -> std::io::Result<TempDir> {
+        use std::path::PathBuf;
+
+        let mut src = PathBuf::new();
+        src.push(env!("CARGO_MANIFEST_DIR"));
+        src.push("../../test_data");
+        src.push(table_path);
+
+        let test_dir = tempdir().unwrap();
+
+        copy_dir_all(src.as_path(), test_dir.path())?;
+
+        Ok(test_dir)
+    }
+
+    #[tokio::test]
+    async fn test_v0_7_5_migration() {
+        // We migrate to add Fragment.physical_rows and DeletionFile.num_deletions
+        // after this version.
+
+        // Copy over table
+        let test_dir = copy_test_data_to_tmp("v0.7.5/with_deletions").unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+
+        // Assert num rows, deletions, and physical rows are all correct.
+        let dataset = Dataset::open(test_uri).await.unwrap();
+        assert_eq!(dataset.count_rows().await.unwrap(), 90);
+        assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10);
+        let total_physical_rows = futures::stream::iter(dataset.get_fragments())
+            .then(|f| async move { f.physical_rows().await })
+            .try_fold(0, |acc, x| async move { Ok(acc + x) })
+            .await
+            .unwrap();
+        assert_eq!(total_physical_rows, 100);
+
+        // Append 5 rows
+        let schema = Arc::new(ArrowSchema::from(dataset.schema()));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int64Array::from_iter_values(100..105))],
+        )
+        .unwrap();
+        let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
+        let write_params = WriteParams {
+            mode: WriteMode::Append,
+            ..Default::default()
+        };
+        let dataset = Dataset::write(batches, test_uri, Some(write_params))
+            .await
+            .unwrap();
+
+        // Assert num rows, deletions, and physical rows are all correct.
+        assert_eq!(dataset.count_rows().await.unwrap(), 95);
+        assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10);
+        let total_physical_rows = futures::stream::iter(dataset.get_fragments())
+            .then(|f| async move { f.physical_rows().await })
+            .try_fold(0, |acc, x| async move { Ok(acc + x) })
+            .await
+            .unwrap();
+        assert_eq!(total_physical_rows, 105);
+
+        dataset.validate().await.unwrap();
+
+        // Scan data and assert it is as expected.
+        let expected = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int64Array::from_iter_values(
+                (0..10).chain(20..105),
+            ))],
+        )
+        .unwrap();
+        let actual_batches = dataset
+            .scan()
+            .try_into_stream()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap();
+        assert_eq!(actual, expected);
+    }
+
+    #[tokio::test]
+    async fn test_fix_v0_8_0_broken_migration() {
+        // The migration from v0.7.5 was broken in 0.8.0. This validates we can
+        // automatically fix tables that have this problem.
+
+        // Copy over table
+        let test_dir = copy_test_data_to_tmp("v0.8.0/migrated_from_v0.7.5").unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+
+        // Assert num rows, deletions, and physical rows are all correct, even
+        // though stats are bad.
+        let dataset = Dataset::open(test_uri).await.unwrap();
+        assert_eq!(dataset.count_rows().await.unwrap(), 92);
+        assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10);
+        let total_physical_rows = futures::stream::iter(dataset.get_fragments())
+            .then(|f| async move { f.physical_rows().await })
+            .try_fold(0, |acc, x| async move { Ok(acc + x) })
+            .await
+            .unwrap();
+        assert_eq!(total_physical_rows, 102);
+
+        // Append 5 rows to table.
+        let schema = Arc::new(ArrowSchema::from(dataset.schema()));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int64Array::from_iter_values(100..105))],
+        )
+        .unwrap();
+        let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
+        let write_params = WriteParams {
+            mode: WriteMode::Append,
+            ..Default::default()
+        };
+        let dataset = Dataset::write(batches, test_uri, Some(write_params))
+            .await
+            .unwrap();
+
+        // Assert statistics are all now correct.
+        let physical_rows: Vec<_> = dataset
+            .get_fragments()
+            .iter()
+            .map(|f| f.metadata.physical_rows)
+            .collect();
+        assert_eq!(physical_rows, vec![Some(100), Some(2), Some(5)]);
+        let num_deletions: Vec<_> = dataset
+            .get_fragments()
+            .iter()
+            .map(|f| {
+                f.metadata
+                    .deletion_file
+                    .as_ref()
+                    .and_then(|df| df.num_deleted_rows)
+            })
+            .collect();
+        assert_eq!(num_deletions, vec![Some(10), None, None]);
+        assert_eq!(dataset.count_rows().await.unwrap(), 97);
+
+        // Scan data and assert it is as expected.
+        let expected = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int64Array::from_iter_values(
+                (0..10).chain(20..100).chain(0..2).chain(100..105),
+            ))],
+        )
+        .unwrap();
+        let actual_batches = dataset
+            .scan()
+            .try_into_stream()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap();
+        assert_eq!(actual, expected);
+    }
 }