feat: support tracking newly inserted and updated rows between versions (#4741)

Co-authored-by: Jack Ye <yezhaoqin@gmail.com>

Author: yanghua

java/lance-jni/Cargo.lock

@@ -633,6 +633,8 @@ impl FromJObjectWithEnv<Fragment> for JObject<'_> {
             deletion_file,
             physical_rows: Some(physical_rows),
             row_id_meta,
+            created_at_version_meta: None,
+            last_updated_at_version_meta: None,
         })
     }
 }

java/lance-jni/src/fragment.rs

@@ -95,4 +95,19 @@ message EncodedU64Array {
         U32Array u32_array = 2;
         U64Array u64_array = 3;
     }
-}
+}
+
+/// A sequence of dataset versions. Similar to RowIdSequence but tracks
+/// version runs. It uses RLE (Run-Length Encoding) to efficiently
+// represent consecutive rows with the same version.
+message RowDatasetVersionSequence {
+    repeated RowDatasetVersionRun runs = 1;
+}
+
+/// A run of rows with the same version.
+message RowDatasetVersionRun {
+    /// The number of consecutive rows with the same version.
+    U64Segment span = 1;
+
+    uint64 version = 2;
+}

protos/rowids.proto

@@ -277,6 +277,20 @@ message DataFragment {
     ExternalFile external_row_ids = 6;
   } // row_id_sequence
 
+  oneof last_updated_at_version_sequence {
+    // If small (< 200KB), the row latest updated versions are stored inline.
+    bytes inline_last_updated_at_versions = 7;
+    // Otherwise, stored as part of a file.
+    ExternalFile external_last_updated_at_versions = 8;
+  } // last_updated_at_version_sequence
+
+  oneof created_at_version_sequence {
+    // If small (< 200KB), the row created at versions are stored inline.
+    bytes inline_created_at_versions = 9;
+    // Otherwise, stored as part of a file.
+    ExternalFile external_created_at_versions = 10;
+  } // created_at_version_sequence
+
   // Number of original rows in the fragment, this includes rows that are now marked with
   // deletion tombstones. To compute the current number of rows, subtract
   // `deletion_file.num_deleted_rows` from this value.

protos/table.proto

@@ -27,6 +27,9 @@
 from .lance import (
     DeletionFile as DeletionFile,
 )
+from .lance import (
+    RowDatasetVersionMeta as RowDatasetVersionMeta,
+)
 from .lance import (
     RowIdMeta as RowIdMeta,
 )
@@ -67,13 +70,19 @@ class FragmentMetadata:
         The deletion file, if any.
     row_id_meta : Optional[RowIdMeta]
         The row id metadata, if any.
+    created_at_version_meta : Optional[RowDatasetVersionMeta]
+        The row created at version metadata, if any.
+    last_updated_at_version_meta : Optional[RowDatasetVersionMeta]
+        The row last updated at version metadata, if any.
     """
 
     id: int
     files: List[DataFile]
     physical_rows: int
     deletion_file: Optional[DeletionFile] = None
     row_id_meta: Optional[RowIdMeta] = None
+    created_at_version_meta: Optional[RowDatasetVersionMeta] = None
+    last_updated_at_version_meta: Optional[RowDatasetVersionMeta] = None
 
     @property
     def num_deletions(self) -> int:
@@ -110,6 +119,16 @@ def to_json(self) -> dict:
             row_id_meta=(
                 self.row_id_meta.asdict() if self.row_id_meta is not None else None
             ),
+            created_at_version_meta=(
+                json.loads(self.created_at_version_meta.json())
+                if self.created_at_version_meta is not None
+                else None
+            ),
+            last_updated_at_version_meta=(
+                json.loads(self.last_updated_at_version_meta.json())
+                if self.last_updated_at_version_meta is not None
+                else None
+            ),
         )
 
     @staticmethod
@@ -124,12 +143,26 @@ def from_json(json_data: str) -> FragmentMetadata:
         if row_id_meta is not None:
             row_id_meta = RowIdMeta(**row_id_meta)
 
+        created_at_version_meta = json_data.get("created_at_version_meta")
+        if created_at_version_meta is not None:
+            created_at_version_meta = RowDatasetVersionMeta.from_json(
+                json.dumps(created_at_version_meta)
+            )
+
+        last_updated_at_version_meta = json_data.get("last_updated_at_version_meta")
+        if last_updated_at_version_meta is not None:
+            last_updated_at_version_meta = RowDatasetVersionMeta.from_json(
+                json.dumps(last_updated_at_version_meta)
+            )
+
         return FragmentMetadata(
             id=json_data["id"],
             files=[DataFile(**f) for f in json_data["files"]],
             physical_rows=json_data["physical_rows"],
             deletion_file=deletion_file,
             row_id_meta=row_id_meta,
+            created_at_version_meta=created_at_version_meta,
+            last_updated_at_version_meta=last_updated_at_version_meta,
         )
 
 

python/python/lance/fragment.py

@@ -278,7 +278,8 @@ def test_fragment_meta():
         "column_indices=[], file_major_version=0, file_minor_version=0, "
         "file_size_bytes=100), DataFile(path='1.lance', fields=[1], column_indices=[], "
         "file_major_version=0, file_minor_version=0, file_size_bytes=None)], "
-        "physical_rows=100, deletion_file=None, row_id_meta=None)"
+        "physical_rows=100, deletion_file=None, row_id_meta=None, "
+        "created_at_version_meta=None, last_updated_at_version_meta=None)"
     )
 
 

python/python/tests/test_fragment.py

@@ -26,7 +26,9 @@ use lance::dataset::transaction::{Operation, Transaction};
 use lance::dataset::{InsertBuilder, NewColumnTransform};
 use lance::Error;
 use lance_io::utils::CachedFileSize;
-use lance_table::format::{DataFile, DeletionFile, DeletionFileType, Fragment, RowIdMeta};
+use lance_table::format::{
+    DataFile, DeletionFile, DeletionFileType, Fragment, RowDatasetVersionMeta, RowIdMeta,
+};
 use lance_table::io::deletion::deletion_file_path;
 use object_store::path::Path;
 use pyo3::basic::CompareOp;
@@ -593,6 +595,9 @@ impl PyDeletionFile {
 #[pyclass(name = "RowIdMeta", module = "lance.fragment")]
 pub struct PyRowIdMeta(pub RowIdMeta);
 
+#[pyclass(name = "RowDatasetVersionMeta", module = "lance.fragment")]
+pub struct PyRowDatasetVersionMeta(pub RowDatasetVersionMeta);
+
 #[pymethods]
 impl PyRowIdMeta {
     fn asdict(&self) -> PyResult<Bound<'_, PyDict>> {
@@ -639,6 +644,55 @@ impl PyRowIdMeta {
     }
 }
 
+#[pymethods]
+impl PyRowDatasetVersionMeta {
+    fn asdict(&self) -> PyResult<Bound<'_, PyDict>> {
+        Err(PyNotImplementedError::new_err(
+            "PyRowDatasetVersionMeta.asdict is not yet supported.",
+        ))
+    }
+
+    pub fn json(&self) -> PyResult<String> {
+        serde_json::to_string(&self.0).map_err(|err| {
+            PyValueError::new_err(format!(
+                "Could not serialize RowDatasetVersionMeta due to error: {}",
+                err
+            ))
+        })
+    }
+
+    #[staticmethod]
+    pub fn from_json(json: String) -> PyResult<Self> {
+        let dataset_version_meta = serde_json::from_str(&json).map_err(|err| {
+            PyValueError::new_err(format!(
+                "Could not load RowDatasetVersionMeta due to error: {}",
+                err
+            ))
+        })?;
+        Ok(Self(dataset_version_meta))
+    }
+
+    fn __reduce__(&self, py: Python<'_>) -> PyResult<(PyObject, PyObject)> {
+        let state = self.json()?;
+        let state = PyTuple::new(py, vec![state])?.extract()?;
+        let from_json = PyModule::import(py, "lance.fragment")?
+            .getattr("RowDatasetVersionMeta")?
+            .getattr("from_json")?
+            .extract()?;
+        Ok((from_json, state))
+    }
+
+    pub fn __richcmp__(&self, other: PyRef<'_, Self>, op: CompareOp) -> PyResult<bool> {
+        match op {
+            CompareOp::Eq => Ok(self.0 == other.0),
+            CompareOp::Ne => Ok(self.0 != other.0),
+            _ => Err(PyNotImplementedError::new_err(
+                "Only == and != are supported for RowDatasetVersionMeta",
+            )),
+        }
+    }
+}
+
 #[pyclass(name = "FragmentSession", module = "_lib", subclass)]
 #[derive(Clone)]
 pub struct FragmentSession {
@@ -670,13 +724,21 @@ impl FromPyObject<'_> for PyLance<Fragment> {
 
         let row_id_meta: Option<PyRef<PyRowIdMeta>> = ob.getattr("row_id_meta")?.extract()?;
         let row_id_meta = row_id_meta.map(|r| r.0.clone());
+        let last_updated_at_version_meta: Option<PyRef<PyRowDatasetVersionMeta>> =
+            ob.getattr("last_updated_at_version_meta")?.extract()?;
+        let last_updated_at_version_meta = last_updated_at_version_meta.map(|r| r.0.clone());
+        let created_at_version_meta: Option<PyRef<PyRowDatasetVersionMeta>> =
+            ob.getattr("created_at_version_meta")?.extract()?;
+        let created_at_version_meta = created_at_version_meta.map(|r| r.0.clone());
 
         Ok(Self(Fragment {
             id: ob.getattr("id")?.extract()?,
             files,
             deletion_file,
             physical_rows: ob.getattr("physical_rows")?.extract()?,
             row_id_meta,
+            last_updated_at_version_meta,
+            created_at_version_meta,
         }))
     }
 }
@@ -699,13 +761,25 @@ impl<'py> IntoPyObject<'py> for PyLance<&Fragment> {
             .as_ref()
             .map(|f| PyDeletionFile(f.clone()));
         let row_id_meta = self.0.row_id_meta.as_ref().map(|r| PyRowIdMeta(r.clone()));
+        let last_updated_at_version_meta = self
+            .0
+            .last_updated_at_version_meta
+            .as_ref()
+            .map(|r| PyRowDatasetVersionMeta(r.clone()));
+        let created_at_version_meta = self
+            .0
+            .created_at_version_meta
+            .as_ref()
+            .map(|r| PyRowDatasetVersionMeta(r.clone()));
 
         cls.call1((
             self.0.id,
             files,
             self.0.physical_rows,
             deletion_file,
             row_id_meta,
+            created_at_version_meta,
+            last_updated_at_version_meta,
         ))
     }
 }

python/src/fragment.rs

@@ -89,7 +89,7 @@ use crate::utils::Hnsw;
 use crate::utils::KMeans;
 pub use dataset::write_dataset;
 pub use dataset::Dataset;
-use fragment::{FileFragment, PyDeletionFile, PyRowIdMeta};
+use fragment::{FileFragment, PyDeletionFile, PyRowDatasetVersionMeta, PyRowIdMeta};
 pub use indices::register_indices;
 pub use reader::LanceReader;
 pub use scanner::Scanner;
@@ -244,6 +244,7 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<FileFragment>()?;
     m.add_class::<PyDeletionFile>()?;
     m.add_class::<PyRowIdMeta>()?;
+    m.add_class::<PyRowDatasetVersionMeta>()?;
     m.add_class::<MergeInsertBuilder>()?;
     m.add_class::<LanceBlobFile>()?;
     m.add_class::<LanceFileReader>()?;

python/src/lib.rs

@@ -38,8 +38,8 @@ pub struct LanceReader {
 }
 
 impl LanceReader {
-    pub async fn try_new(scanner: Arc<LanceScanner>) -> ::lance::error::Result<Self> {
-        let stream = scanner.try_into_stream().await?;
+    pub async fn try_new(mut scanner: Arc<LanceScanner>) -> ::lance::error::Result<Self> {
+        let stream = Arc::make_mut(&mut scanner).try_into_stream().await?;
         let schema = stream.schema();
         Ok(Self {
             schema,