From ea1d2b172ae240c22a89d05b1c66ca6d3a5f90a4 Mon Sep 17 00:00:00 2001 From: XY Zhan Date: Fri, 16 Oct 2026 15:27:53 -0400 Subject: [PATCH] fix: require a primary key to enable MemWAL (#4561) ## Why Every MemWAL write and read resolves rows by the table's unenforced primary key, but `enable_memwal` accepted a table without one. The spec installed, and every MemWAL `merge_insert` then failed with "table has a MemWAL index but no unenforced primary key" until MemWAL was disabled. Lance's MemWAL initialization doesn't catch it: its key checks pass when there is no key. ## What changes `enable_memwal` refuses a spec on a table with no unenforced primary key, for every sharding mode, before anything is committed. The check runs against the dataset the spec is committed to; a key cannot be removed once set, and an overwrite or restore of the table conflicts with that commit. Servers that install specs through this method get the check too. The Rust doc on `Table::enable_memwal` now states the requirement; the Python and TypeScript docs already did. ## Tests - New `test_enable_memwal_requires_primary_key` (Rust and Python): unsharded, bucket and identity specs are all refused without a key, a refused spec leaves nothing installed, and the same table accepts a spec once a key is set. The Rust test fails with the check disabled. - Tests that installed a spec on a table without a key now declare one first (6 Rust, 3 Python). - `cargo nextest run -p lancedb --features remote --lib`: 1529 passed. `cargo clippy -p lancedb --features remote --tests --examples -- -D warnings` and `cargo fmt --check` are clean; no integration test or example installs a spec. - Python: `test_memwal_spec.py` and `test_merge_insert_lsm.py`, 43 passed; `ruff format` and `ruff check` are clean. --- python/python/tests/test_memwal_spec.py | 15 +++++++ rust/lancedb/src/table.rs | 59 +++++++++++++++++++++++++ rust/lancedb/src/table/merge/lsm.rs | 11 +++++ 3 files changed, 85 insertions(+) diff --git a/python/python/tests/test_memwal_spec.py b/python/python/tests/test_memwal_spec.py index a1ec6a90d..d292c9185 100644 --- a/python/python/tests/test_memwal_spec.py +++ b/python/python/tests/test_memwal_spec.py @@ -45,6 +45,7 @@ def _make_table(tmp_path): def test_enable_memwal_validates(tmp_path): _db, table = _make_table(tmp_path) + table.set_unenforced_primary_key("id") # Out-of-range num_buckets. with pytest.raises(Exception, match="num_buckets"): @@ -60,6 +61,7 @@ def test_enable_memwal_validates(tmp_path): def test_disable_memwal(tmp_path): _db, table = _make_table(tmp_path) + table.set_unenforced_primary_key("id") # unset errors when no spec is set. with pytest.raises(Exception, match="no MemWAL write spec"): @@ -74,6 +76,18 @@ def test_disable_memwal(tmp_path): table.enable_memwal(MemWalWriteSpec.bucket("id", 8)) +def test_enable_memwal_requires_primary_key(tmp_path): + _db, table = _make_table(tmp_path) + for spec in ( + MemWalWriteSpec.unsharded(), + MemWalWriteSpec.bucket("id", 4), + MemWalWriteSpec.identity("id"), + ): + with pytest.raises(Exception, match="primary key"): + table.enable_memwal(spec) + assert table.get_memwal_spec() is None + + def test_set_unsharded_spec(tmp_path): _db, table = _make_table(tmp_path) # Lance MemWAL still requires a primary key on the dataset; Unsharded @@ -206,6 +220,7 @@ async def test_async_get_memwal_spec(tmp_path): ) assert await table.get_memwal_spec() is None + await table.set_unenforced_primary_key("id") # A real scalar index is needed to name it as a maintained index. await table.create_index("id", config=BTree()) diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index c0a4a719e..c6193d750 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -2015,6 +2015,9 @@ impl Table { /// - [`MemWalWriteSpec::identity`] — shard by the raw value of a scalar column. /// - [`MemWalWriteSpec::unsharded`] — route every write to a single shard. /// + /// Every spec requires the table to declare an unenforced primary key + /// ([`Table::set_unenforced_primary_key`]); otherwise this fails. + /// /// A table carries one spec: this fails while one is installed, since the /// SSTables already written were homed under it. /// [`Table::disable_memwal`] removes one. @@ -5541,6 +5544,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); // Reject num_buckets out of range. for bad in [0u32, 1025] { @@ -5585,6 +5589,56 @@ mod tests { assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err); } + /// A spec is refused on a table without an unenforced primary key, for every + /// sharding mode, and a refused spec leaves nothing installed. + #[tokio::test] + async fn test_enable_memwal_requires_primary_key() { + use lance::dataset::mem_wal::DatasetMemWalExt; + + let tmp_dir = tempdir().unwrap(); + let conn = ConnectBuilder::new(tmp_dir.path().to_str().unwrap()) + .read_consistency_interval(Duration::from_secs(0)) + .execute() + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::Int64Array::from(vec![1, 2]))], + ) + .unwrap(); + let reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)); + let table = conn.create_table("t", reader).execute().await.unwrap(); + let installed = |table: &Table| { + let table = table.clone(); + async move { + let dataset = table.as_native().unwrap().dataset.get().await.unwrap(); + dataset.mem_wal_index_details().await.unwrap().is_some() + } + }; + + for spec in [ + MemWalWriteSpec::unsharded(), + MemWalWriteSpec::bucket("id", 4), + MemWalWriteSpec::identity("id"), + ] { + let err = table.enable_memwal(spec).await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("primary key")), + "got {err:?}" + ); + assert!(!installed(&table).await); + } + + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .enable_memwal(MemWalWriteSpec::unsharded()) + .await + .unwrap(); + assert!(installed(&table).await); + } + #[tokio::test] async fn test_enable_memwal_unsharded() { use lance::dataset::mem_wal::DatasetMemWalExt; @@ -5606,6 +5660,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); table .enable_memwal(MemWalWriteSpec::unsharded()) @@ -5652,6 +5707,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); table .enable_memwal( @@ -5703,6 +5759,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); // unset errors when no spec is set. table.disable_memwal().await.unwrap_err(); @@ -5760,6 +5817,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); // No spec installed yet. assert_eq!(table.get_memwal_spec().await.unwrap(), None); @@ -5826,6 +5884,7 @@ mod tests { .await .unwrap(); let table = conn.create_table("t", reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); table .create_index(&["id"], Index::BTree(Default::default())) diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index d26be28cc..39e4fbf28 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -127,6 +127,17 @@ pub(crate) async fn enable_memwal(table: &NativeTable, spec: MemWalWriteSpec) -> .into(), }); } + // Every MemWAL write and read resolves rows by the unenforced primary key, + // so a spec on a table without one leaves the table unwritable through the + // MemWAL path. A key cannot be removed once set, and an overwrite or restore + // of the table conflicts with the commit below. + if dataset.schema().unenforced_primary_key().is_empty() { + return Err(Error::InvalidInput { + message: "a MemWAL write spec requires the table to declare an unenforced \ + primary key; set one with set_unenforced_primary_key first" + .into(), + }); + } let mut builder = dataset .initialize_mem_wal() .mem_index_registry(table.mem_index_registry());