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());