Skip to content

Commit c3dd12b

Browse files
committed
fix: support JSON serialization of FragmentMetadata with stable row IDs
`FragmentMetadata.to_json()` raised `NotImplementedError` when `row_id_meta` was present (i.e. `enable_stable_row_ids=True`), because `PyRowIdMeta.asdict()` was not implemented. Implement `asdict()` using `pythonize` to convert the Rust enum directly to a Python dict. Also fix `from_json()` to deserialize `row_id_meta` via `RowIdMeta.from_json()` instead of attempting `RowIdMeta(**dict)`.
1 parent 2c20d75 commit c3dd12b

File tree

3 files changed

+14
-7
lines changed

3 files changed

+14
-7
lines changed

python/python/lance/fragment.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ def from_json(json_data: str) -> FragmentMetadata:
143143

144144
row_id_meta = json_data.get("row_id_meta")
145145
if row_id_meta is not None:
146-
row_id_meta = RowIdMeta(**row_id_meta)
146+
row_id_meta = RowIdMeta.from_json(json.dumps(row_id_meta))
147147

148148
created_at_version_meta = json_data.get("created_at_version_meta")
149149
if created_at_version_meta is not None:

python/python/tests/test_fragment.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -470,11 +470,18 @@ def test_fragment_metadata_pickle(tmp_path: Path, enable_stable_row_ids: bool):
470470
if enable_stable_row_ids:
471471
assert frag_meta.row_id_meta is not None
472472

473-
# Pickle and unpickle the fragment metadata
473+
# Pickle round-trip
474474
round_trip = pickle.loads(pickle.dumps(frag_meta))
475-
476475
assert frag_meta == round_trip
477476

477+
# JSON round-trip
478+
json_data = frag_meta.to_json()
479+
json_round_trip = FragmentMetadata.from_json(json.dumps(json_data))
480+
assert frag_meta.id == json_round_trip.id
481+
assert frag_meta.physical_rows == json_round_trip.physical_rows
482+
if enable_stable_row_ids:
483+
assert json_round_trip.row_id_meta is not None
484+
478485

479486
def test_deletion_file_with_base_id_serialization():
480487
"""Test that DeletionFile with base_id serializes correctly."""

python/src/fragment.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -627,10 +627,10 @@ pub struct PyRowDatasetVersionMeta(pub RowDatasetVersionMeta);
627627

628628
#[pymethods]
629629
impl PyRowIdMeta {
630-
fn asdict(&self) -> PyResult<Bound<'_, PyDict>> {
631-
Err(PyNotImplementedError::new_err(
632-
"PyRowIdMeta.asdict is not yet supported.s",
633-
))
630+
fn asdict(&self, py: Python<'_>) -> PyResult<PyObject> {
631+
pythonize::pythonize(py, &self.0)
632+
.map(|b| b.unbind())
633+
.map_err(|err| PyValueError::new_err(format!("Could not convert RowIdMeta: {}", err)))
634634
}
635635

636636
pub fn json(&self) -> PyResult<String> {

0 commit comments

Comments
 (0)