feat: add support for pickling fragment metadata (#3497)

Author: westonpace

python/python/lance/lance/fragment.pyi

@@ -56,5 +56,65 @@ class DeletionFile:
         """
         ...
 
+    def json(self) -> str:
+        """Get a JSON representation of the deletion file.
+
+        Returns
+        -------
+        str
+
+        Warning
+        -------
+        The JSON representation is not guaranteed to be stable across versions.
+        """
+        ...
+
+    @classmethod
+    def from_json(json: str) -> DeletionFile:
+        """
+        Load a deletion file from a JSON representation.
+
+        Parameters
+        ----------
+        json : str
+            The JSON representation of the deletion file.
+
+        Returns
+        -------
+        DeletionFile
+        """
+        ...
+
+    def __reduce__(self) -> tuple: ...
+
 class RowIdMeta:
-    pass
+    def json(self) -> str:
+        """Get a JSON representation of the row id metadata.
+
+        Returns
+        -------
+        str
+
+        Warning
+        -------
+        The JSON representation is not guaranteed to be stable across versions.
+        """
+        ...
+
+    @classmethod
+    def from_json(json: str) -> RowIdMeta:
+        """
+        Load row id metadata from a JSON representation.
+
+        Parameters
+        ----------
+        json : str
+            The JSON representation of the row id metadata.
+
+        Returns
+        -------
+        RowIdMeta
+        """
+        ...
+
+    def __reduce__(self) -> tuple: ...

python/python/tests/test_fragment.py

@@ -3,6 +3,7 @@
 
 import json
 import multiprocessing
+import pickle
 import uuid
 from pathlib import Path
 
@@ -435,3 +436,26 @@ def test_fragment_count_rows(tmp_path: Path):
     assert fragments[0].count_rows() == 800
     assert fragments[0].count_rows("a < 200") == 200
     assert fragments[0].count_rows(pc.field("a") < 200) == 200
+
+
+@pytest.mark.parametrize("enable_move_stable_row_ids", [False, True])
+def test_fragment_metadata_pickle(tmp_path: Path, enable_move_stable_row_ids: bool):
+    ds = write_dataset(
+        pa.table({"a": range(100)}),
+        tmp_path,
+        enable_move_stable_row_ids=enable_move_stable_row_ids,
+    )
+    # Create a deletion file
+    ds.delete("a < 50")
+    fragment = ds.get_fragments()[0]
+
+    frag_meta = fragment.metadata
+
+    assert frag_meta.deletion_file is not None
+    if enable_move_stable_row_ids:
+        assert frag_meta.row_id_meta is not None
+
+    # Pickle and unpickle the fragment metadata
+    round_trip = pickle.loads(pickle.dumps(frag_meta))
+
+    assert frag_meta == round_trip

python/src/fragment.rs

@@ -27,6 +27,8 @@ use lance::Error;
 use lance_table::format::{DataFile, DeletionFile, DeletionFileType, Fragment, RowIdMeta};
 use lance_table::io::deletion::deletion_file_path;
 use object_store::path::Path;
+use pyo3::basic::CompareOp;
+use pyo3::types::PyTuple;
 use pyo3::{exceptions::*, types::PyDict};
 use pyo3::{intern, prelude::*};
 use snafu::location;
@@ -507,6 +509,43 @@ impl PyDeletionFile {
         };
         Ok(deletion_file_path(&base_path, fragment_id, &self.0).to_string())
     }
+
+    pub fn json(&self) -> PyResult<String> {
+        serde_json::to_string(&self.0).map_err(|err| {
+            PyValueError::new_err(format!(
+                "Could not dump CompactionPlan due to error: {}",
+                err
+            ))
+        })
+    }
+
+    #[staticmethod]
+    pub fn from_json(json: String) -> PyResult<Self> {
+        let deletion_file = serde_json::from_str(&json).map_err(|err| {
+            PyValueError::new_err(format!("Could not load DeletionFile due to error: {}", err))
+        })?;
+        Ok(Self(deletion_file))
+    }
+
+    fn __reduce__(&self, py: Python<'_>) -> PyResult<(PyObject, PyObject)> {
+        let state = self.json()?;
+        let state = PyTuple::new_bound(py, vec![state]).extract()?;
+        let from_json = PyModule::import_bound(py, "lance.fragment")?
+            .getattr("DeletionFile")?
+            .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 CompactionTask",
+            )),
+        }
+    }
 }
 
 #[pyclass(name = "RowIdMeta", module = "lance.fragment")]
@@ -519,6 +558,43 @@ impl PyRowIdMeta {
             "PyRowIdMeta.asdict is not yet supported.s",
         ))
     }
+
+    pub fn json(&self) -> PyResult<String> {
+        serde_json::to_string(&self.0).map_err(|err| {
+            PyValueError::new_err(format!(
+                "Could not dump CompactionPlan due to error: {}",
+                err
+            ))
+        })
+    }
+
+    #[staticmethod]
+    pub fn from_json(json: String) -> PyResult<Self> {
+        let row_id_meta = serde_json::from_str(&json).map_err(|err| {
+            PyValueError::new_err(format!("Could not load RowIdMeta due to error: {}", err))
+        })?;
+        Ok(Self(row_id_meta))
+    }
+
+    fn __reduce__(&self, py: Python<'_>) -> PyResult<(PyObject, PyObject)> {
+        let state = self.json()?;
+        let state = PyTuple::new_bound(py, vec![state]).extract()?;
+        let from_json = PyModule::import_bound(py, "lance.fragment")?
+            .getattr("RowIdMeta")?
+            .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 CompactionTask",
+            )),
+        }
+    }
 }
 
 impl FromPyObject<'_> for PyLance<Fragment> {