From 979ac020cc8d09401d8ff3676cb042ca3019d0c5 Mon Sep 17 00:00:00 2001 From: kid Date: Tue, 7 Jul 2026 11:33:11 +0800 Subject: [PATCH] feat(python): expose MemWAL shard delete --- python/python/lance/lance/__init__.pyi | 1 + python/python/lance/mem_wal.py | 31 ++++++++++++++++++++ python/python/tests/test_mem_wal.py | 28 ++++++++++++++++++ python/src/mem_wal.rs | 39 ++++++++++++++++++++++---- 4 files changed, 94 insertions(+), 5 deletions(-) diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 53707bda41b..d131c08e7c8 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -712,6 +712,7 @@ class _ShardSnapshot: class _ShardWriter: shard_id: str def put(self, data: Any) -> None: ... + def delete(self, keys: Any) -> None: ... def close(self) -> None: ... def stats(self) -> Dict[str, Any]: ... def memtable_stats(self) -> Dict[str, Any]: ... diff --git a/python/python/lance/mem_wal.py b/python/python/lance/mem_wal.py index f87e811f830..426a5f7ee50 100644 --- a/python/python/lance/mem_wal.py +++ b/python/python/lance/mem_wal.py @@ -187,6 +187,7 @@ class ShardWriter: with dataset.mem_wal_writer(shard_id) as writer: writer.put(batch) + writer.delete(pa.table({"id": [1]})) Parameters ---------- @@ -224,6 +225,36 @@ def put(self, data, *, schema: Optional[pa.Schema] = None) -> None: reader = _coerce_reader(data, schema) self._raw.put(reader) + def delete(self, keys, *, schema: Optional[pa.Schema] = None) -> None: + """Delete rows by primary key from the MemWAL. + + Parameters + ---------- + keys : ReaderLike + Any Arrow-compatible data containing this shard's primary key + column(s). Non-primary-key columns, if present, are ignored by + the Rust core delete path. + schema : pa.Schema, optional + Schema hint, needed when *keys* is a generator. + + Raises + ------ + IOError + If delete validation fails, WAL flush fails, or the writer has + already been closed. Delete validation is centralized in Rust and + includes checks for primary-key metadata and tombstone-compatible + nullable non-key columns. + + Examples + -------- + :: + + with dataset.mem_wal_writer(shard_id) as writer: + writer.delete(pa.table({"id": [42]})) + """ + reader = _coerce_reader(keys, schema) + self._raw.delete(reader) + def close(self) -> None: """Flush and close the writer. diff --git a/python/python/tests/test_mem_wal.py b/python/python/tests/test_mem_wal.py index 95596a7123e..2871baad986 100644 --- a/python/python/tests/test_mem_wal.py +++ b/python/python/tests/test_mem_wal.py @@ -196,6 +196,34 @@ def test_shard_writer_lsm_scanner_includes_own_flushed_generations(tmp_path): time.sleep(0.05) +def test_shard_writer_delete_binding_masks_base_row(tmp_path): + ds_path = str(tmp_path / "base") + shard_id = str(uuid.uuid4()) + ds = lance.write_dataset( + _lookup_table([1, 2, 3], "base"), ds_path, schema=_LOOKUP_SCHEMA + ) + ds.initialize_mem_wal() + + delete_keys = pa.table({"id": pa.array([2], type=pa.int64())}) + + with ds.mem_wal_writer( + shard_id, + durable_write=True, + sync_indexed_write=True, + max_wal_buffer_size=1, + max_wal_flush_interval_ms=10, + ) as writer: + writer.put(_lookup_table([4], "writer")) + writer.delete(delete_keys) + table = writer.lsm_scanner().to_table() + + rows = {row["id"]: row["name"] for row in table.to_pylist()} + assert rows[1] == "base_1" + assert 2 not in rows, "deleted base row should be masked by the tombstone" + assert rows[3] == "base_3" + assert rows[4] == "writer_4" + + _VDIM = 4 # matches Rust test fixture dimension diff --git a/python/src/mem_wal.rs b/python/src/mem_wal.rs index 812c5f42f22..e415acfc38d 100644 --- a/python/src/mem_wal.rs +++ b/python/src/mem_wal.rs @@ -238,6 +238,14 @@ struct ClosedShardWriterState { memtable_stats: MemTableStats, } +fn collect_record_batches(data: &Bound<'_, PyAny>) -> PyResult> { + let reader = ArrowArrayStreamReader::from_pyarrow_bound(data) + .map_err(|e| PyValueError::new_err(format!("Cannot read data as Arrow: {}", e)))?; + reader + .collect::>() + .map_err(|e| PyIOError::new_err(format!("Failed to read batches: {}", e))) +} + #[pymethods] impl PyShardWriter { /// Write data batches to the MemWAL. @@ -245,11 +253,7 @@ impl PyShardWriter { /// Accepts any PyArrow-compatible data source (RecordBatch, Table, /// or an Arrow stream reader). pub fn put(&self, py: Python<'_>, data: &Bound<'_, PyAny>) -> PyResult<()> { - let reader = ArrowArrayStreamReader::from_pyarrow_bound(data) - .map_err(|e| PyValueError::new_err(format!("Cannot read data as Arrow: {}", e)))?; - let batches: Vec = reader - .collect::>() - .map_err(|e| PyIOError::new_err(format!("Failed to read batches: {}", e)))?; + let batches = collect_record_batches(data)?; if batches.is_empty() { return Ok(()); @@ -268,6 +272,31 @@ impl PyShardWriter { .map_err(|e: lance::Error| PyIOError::new_err(e.to_string())) } + /// Delete rows from the MemWAL by primary key. + /// + /// Accepts any PyArrow-compatible data source carrying the shard's primary + /// key column(s). Rust core validates that primary keys exist and builds the + /// tombstone rows. + pub fn delete(&self, py: Python<'_>, keys: &Bound<'_, PyAny>) -> PyResult<()> { + let batches = collect_record_batches(keys)?; + + if batches.is_empty() { + return Ok(()); + } + + let inner = self.inner.clone(); + rt().block_on(Some(py), async move { + let guard = inner.lock().await; + match guard.as_ref() { + Some(writer) => writer.delete(batches).await.map(|_| ()), + None => Err(lance_core::Error::invalid_input( + "ShardWriter is already closed", + )), + } + })? + .map_err(|e: lance::Error| PyIOError::new_err(e.to_string())) + } + /// Flush pending data and close the writer. /// /// After close(), calling put() will raise an error.