Skip to content

Commit ef256e2

Browse files
hamersawclaude
andauthored
feat(java): expose MemWAL shard delete (#7688)
## What Mirrors the Python `ShardWriter.delete` binding on the Java side, bringing the Java MemWAL bindings to parity with Python. The Rust core `ShardWriter::delete` and the Python binding already exist; this wires up the Java surface. - **`java/lance-jni/src/mem_wal.rs`** — `nativeDelete` JNI entry point (`inner_delete`) that streams Arrow key batches into `ShardWriter::delete`, mirroring `nativePut`/`inner_put`. - **`ShardWriter.java`** — public `delete(ArrowReader)` wrapper mirroring `put`, with Javadoc copied from the Rust core (tombstone semantics, nullable-column and primary-key constraints). - **`MemWalTest.java`** — integration test `testShardWriterDeleteMasksBaseRow` asserting a tombstone masks the deleted base row, mirroring the Python `test_shard_writer_delete_binding_masks_base_row`. All validation and tombstone construction stay centralized in the Rust core — the binding is a thin wrapper, per the cross-language binding guidelines. ## Testing - `cargo clippy --tests --manifest-path ./lance-jni/Cargo.toml` — clean - `./mvnw spotless:check` + `cargo fmt` — clean - `./mvnw test -Dtest=MemWalTest#testShardWriterDeleteMasksBaseRow` — `Tests run: 1, Failures: 0` 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for deleting records through the shard writer API. * Deletions now work using primary-key-only input and are reflected as hidden tombstone rows. * **Bug Fixes** * Deleted rows are now masked from query results while preserving other existing records and newly added data. * Added validation and safer handling for empty delete inputs and invalid usage. * **Tests** * Added an integration test covering delete behavior and tombstone masking. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent af8a90f commit ef256e2

3 files changed

Lines changed: 114 additions & 0 deletions

File tree

java/lance-jni/src/mem_wal.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,29 @@ fn inner_put(env: &mut JNIEnv, this: JObject, stream_addr: jlong) -> Result<()>
181181
Ok(())
182182
}
183183

184+
#[unsafe(no_mangle)]
185+
pub extern "system" fn Java_org_lance_memwal_ShardWriter_nativeDelete(
186+
mut env: JNIEnv,
187+
this: JObject,
188+
stream_addr: jlong,
189+
) {
190+
ok_or_throw_without_return!(env, inner_delete(&mut env, this, stream_addr));
191+
}
192+
193+
fn inner_delete(env: &mut JNIEnv, this: JObject, stream_addr: jlong) -> Result<()> {
194+
let stream_ptr = stream_addr as *mut FFI_ArrowArrayStream;
195+
let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?;
196+
let batches: Vec<RecordBatch> = reader.collect::<std::result::Result<_, _>>()?;
197+
if batches.is_empty() {
198+
return Ok(());
199+
}
200+
201+
let guard =
202+
unsafe { env.get_rust_field::<_, _, BlockingShardWriter>(&this, NATIVE_SHARD_WRITER) }?;
203+
RT.block_on(guard.writer.delete(batches))?;
204+
Ok(())
205+
}
206+
184207
/// Test-support: write a primary-key dedup sidecar (`_pk_index/`) for a
185208
/// flushed-generation dataset already staged at `gen_path`, mirroring what
186209
/// production flush emits. Lets Java tests stage a *faithful* flushed

java/src/main/java/org/lance/memwal/ShardWriter.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
* <pre>{@code
3737
* try (ShardWriter writer = dataset.memWalWriter(shardId)) {
3838
* writer.put(reader);
39+
* writer.delete(keys);
3940
* }
4041
* }</pre>
4142
*
@@ -99,6 +100,34 @@ public void put(ArrowReader reader) {
99100

100101
private native void nativePut(long streamAddress);
101102

103+
/**
104+
* Delete rows from the MemWAL by primary key.
105+
*
106+
* <p>Each batch in {@code reader} must carry this shard's primary key column(s); other columns
107+
* are ignored. Lance builds a tombstone row per key — the primary key plus {@code _tombstone =
108+
* true} and null in every other column — and appends it like an ordinary write. The tombstone is
109+
* the newest value for its key: it wins newest-per-PK resolution (suppressing the older real row)
110+
* and is then dropped from query results.
111+
*
112+
* <p>Only supported in memtable mode. Because a tombstone nulls every non-PK column, those
113+
* columns must be nullable in the base schema; deleting against a schema with a non-nullable
114+
* non-PK column errors. Deleting on a shard with no primary key columns also errors.
115+
*
116+
* @param reader the keys to delete; consumed fully by this call
117+
*/
118+
public void delete(ArrowReader reader) {
119+
Preconditions.checkNotNull(reader, "reader must not be null");
120+
try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
121+
Preconditions.checkArgument(nativeShardWriterHandle != 0, "ShardWriter is closed");
122+
try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
123+
Data.exportArrayStream(allocator, reader, stream);
124+
nativeDelete(stream.memoryAddress());
125+
}
126+
}
127+
}
128+
129+
private native void nativeDelete(long streamAddress);
130+
102131
/** Return a snapshot of cumulative write statistics. */
103132
public WriteStats stats() {
104133
try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {

java/src/test/java/org/lance/memwal/MemWalTest.java

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,22 @@ private static VectorSchemaRoot lookupRoot(BufferAllocator allocator, long[] ids
9797
return root;
9898
}
9999

100+
/** Build a single-batch root carrying only the {@code id} primary key, for deletes. */
101+
private static VectorSchemaRoot keysRoot(BufferAllocator allocator, long[] ids) {
102+
VectorSchemaRoot root =
103+
VectorSchemaRoot.create(
104+
new Schema(
105+
Collections.singletonList(Field.nullable("id", new ArrowType.Int(64, true)))),
106+
allocator);
107+
BigIntVector idVector = (BigIntVector) root.getVector("id");
108+
idVector.allocateNew(ids.length);
109+
for (int i = 0; i < ids.length; i++) {
110+
idVector.set(i, ids[i]);
111+
}
112+
root.setRowCount(ids.length);
113+
return root;
114+
}
115+
100116
/** Build a single-batch append-only root without primary-key metadata. */
101117
private static VectorSchemaRoot appendOnlyRoot(
102118
BufferAllocator allocator, long[] ids, String prefix) {
@@ -382,6 +398,52 @@ void testShardWriterPutAndLsmScanner(@TempDir Path tempDir) throws Exception {
382398
}
383399
}
384400

401+
@Test
402+
void testShardWriterDeleteMasksBaseRow(@TempDir Path tempDir) throws Exception {
403+
String path = tempDir.resolve("base").toString();
404+
String shardId = UUID.randomUUID().toString();
405+
try (BufferAllocator allocator = new RootAllocator();
406+
Dataset dataset = writeLookupDataset(allocator, path, new long[] {1, 2, 3}, "base")) {
407+
dataset.initializeMemWal(new InitializeMemWalParams());
408+
409+
ShardWriterConfig config =
410+
new ShardWriterConfig()
411+
.withDurableWrite(true)
412+
.withSyncIndexedWrite(true)
413+
.withMaxWalBufferSize(1)
414+
.withMaxWalFlushIntervalMs(10);
415+
416+
try (ShardWriter writer = dataset.memWalWriter(shardId, config)) {
417+
try (VectorSchemaRoot root = lookupRoot(allocator, new long[] {4}, "writer");
418+
ArrowReader reader = toReader(allocator, root)) {
419+
writer.put(reader);
420+
}
421+
try (VectorSchemaRoot keys = keysRoot(allocator, new long[] {2});
422+
ArrowReader reader = toReader(allocator, keys)) {
423+
writer.delete(reader);
424+
}
425+
426+
Map<Long, String> byId = Collections.emptyMap();
427+
long deadline = System.currentTimeMillis() + 10_000;
428+
while (System.currentTimeMillis() < deadline) {
429+
try (LsmScanner scanner = writer.lsmScanner();
430+
ArrowReader reader = scanner.scanBatches()) {
431+
byId = readByName(reader);
432+
}
433+
if (!byId.containsKey(2L) && "writer_4".equals(byId.get(4L))) {
434+
break;
435+
}
436+
Thread.sleep(50);
437+
}
438+
439+
assertEquals("base_1", byId.get(1L));
440+
assertFalse(byId.containsKey(2L), "deleted base row should be masked by the tombstone");
441+
assertEquals("base_3", byId.get(3L));
442+
assertEquals("writer_4", byId.get(4L));
443+
}
444+
}
445+
}
446+
385447
@Test
386448
void testLsmScannerFromSnapshots(@TempDir Path tempDir) throws Exception {
387449
String basePath = tempDir.resolve("base").toString();

0 commit comments

Comments
 (0)