From 6633f83f7bf78350178715a36a422fb13259b841 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Tue, 7 Jul 2026 14:45:33 -0500 Subject: [PATCH] feat: add get_lsm_write_spec to read the installed LSM write spec Adds `Table::get_lsm_write_spec` returning `Option`, the read counterpart to `set_lsm_write_spec` / `unset_lsm_write_spec`. - Rust core: `NativeTable` reconstructs the spec from the MemWAL index details (`mem_wal_index_details`), resolving the shard column via the dataset schema. Read-only, so unaffected by describe_indices filtering. - Remote: `RemoteTable` reads the `__lance_mem_wal` system index via `index/list` with `include_system: true` (kept out of the curated `list_indices` surface) and parses its `details` JSON, using the server-resolved `column` name. - Python and TypeScript bindings, sync + async, with round-trip tests. The remote path depends on lance emitting the resolved shard column name in the MemWAL index details JSON, and on the server honoring `include_system` on `index/list`. Co-Authored-By: Claude Opus 4.8 --- nodejs/__test__/table.test.ts | 40 ++++ nodejs/lancedb/table.ts | 15 ++ nodejs/src/table.rs | 51 +++++ python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/remote/table.py | 4 + python/python/lancedb/table.py | 16 ++ python/python/tests/test_lsm_write_spec.py | 59 ++++++ python/src/table.rs | 14 ++ rust/lancedb/src/remote/table.rs | 214 +++++++++++++++++++++ rust/lancedb/src/table.rs | 98 ++++++++++ rust/lancedb/src/table/merge/lsm.rs | 93 ++++++++- 11 files changed, 604 insertions(+), 1 deletion(-) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index f4ba37361..2ae811842 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2992,6 +2992,46 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => { }), ).rejects.toThrow(); }); + + it("reads back the installed spec via getLsmWriteSpec", async () => { + const conn = await connect(tmpDir.name); + const table = await makeTable(conn); + await table.setUnenforcedPrimaryKey("id"); + + // Nothing installed yet. + expect(await table.getLsmWriteSpec()).toBeUndefined(); + + // Bucket spec round-trips, including writer config defaults. + await table.setLsmWriteSpec({ + specType: "bucket", + column: "id", + numBuckets: 4, + writerConfigDefaults: { durable_write: "false" }, + }); + const spec = await table.getLsmWriteSpec(); + expect(spec).toBeDefined(); + expect(spec?.specType).toBe("bucket"); + expect(spec?.column).toBe("id"); + expect(spec?.numBuckets).toBe(4); + expect(spec?.writerConfigDefaults).toEqual({ durable_write: "false" }); + + // After unset, undefined again. + await table.unsetLsmWriteSpec(); + expect(await table.getLsmWriteSpec()).toBeUndefined(); + + // Identity round-trips (column recovered from the schema). + await table.setLsmWriteSpec({ specType: "identity", column: "id" }); + const identity = await table.getLsmWriteSpec(); + expect(identity?.specType).toBe("identity"); + expect(identity?.column).toBe("id"); + await table.unsetLsmWriteSpec(); + + // Unsharded round-trips (no routing column). + await table.setLsmWriteSpec({ specType: "unsharded" }); + const unsharded = await table.getLsmWriteSpec(); + expect(unsharded?.specType).toBe("unsharded"); + expect(unsharded?.column).toBeFalsy(); + }); }); describe("LSM merge insert", () => { diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index e8d70e9f2..aa6070216 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -585,6 +585,17 @@ export abstract class Table { * @returns {Promise} */ abstract unsetLsmWriteSpec(): Promise; + /** + * Read the {@link LsmWriteSpec} currently installed on this table. + * + * Resolves to `undefined` when the MemWAL LSM write path is not enabled (no + * spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}). + * The returned spec — including its `maintainedIndexes` and + * `writerConfigDefaults` — mirrors what was passed to + * {@link Table#setLsmWriteSpec}. + * @returns {Promise} + */ + abstract getLsmWriteSpec(): Promise; /** * Drain and close any cached MemWAL shard writers held for this table. * @@ -1091,6 +1102,10 @@ export class LocalTable extends Table { return await this.inner.unsetLsmWriteSpec(); } + async getLsmWriteSpec(): Promise { + return (await this.inner.getLsmWriteSpec()) ?? undefined; + } + async closeLsmWriters(): Promise { return await this.inner.closeLsmWriters(); } diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 7558bb4b6..251282d88 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -411,6 +411,16 @@ impl Table { .default_error() } + #[napi(catch_unwind)] + pub async fn get_lsm_write_spec(&self) -> napi::Result> { + let spec = self + .inner_ref()? + .get_lsm_write_spec() + .await + .default_error()?; + Ok(spec.map(LsmWriteSpec::from)) + } + #[napi(catch_unwind)] pub async fn close_lsm_writers(&self) -> napi::Result<()> { self.inner_ref()?.close_lsm_writers().await.default_error() @@ -728,6 +738,47 @@ impl TryFrom for lancedb::table::LsmWriteSpec { } } +impl From for LsmWriteSpec { + fn from(spec: lancedb::table::LsmWriteSpec) -> Self { + use lancedb::table::LsmWriteSpec as Native; + match spec { + Native::Bucket { + column, + num_buckets, + maintained_indexes, + writer_config_defaults, + } => Self { + spec_type: "bucket".to_string(), + column: Some(column), + num_buckets: Some(num_buckets), + maintained_indexes: Some(maintained_indexes), + writer_config_defaults: Some(writer_config_defaults), + }, + Native::Identity { + column, + maintained_indexes, + writer_config_defaults, + } => Self { + spec_type: "identity".to_string(), + column: Some(column), + num_buckets: None, + maintained_indexes: Some(maintained_indexes), + writer_config_defaults: Some(writer_config_defaults), + }, + Native::Unsharded { + maintained_indexes, + writer_config_defaults, + } => Self { + spec_type: "unsharded".to_string(), + column: None, + num_buckets: None, + maintained_indexes: Some(maintained_indexes), + writer_config_defaults: Some(writer_config_defaults), + }, + } + } +} + /// Statistics about a compaction operation. #[napi(object)] #[derive(Clone, Debug)] diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 8ddb28604..8eddb6ef3 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -226,6 +226,7 @@ class Table: async def set_unenforced_primary_key(self, columns: List[str]) -> None: ... async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ... async def unset_lsm_write_spec(self) -> None: ... + async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ... async def close_lsm_writers(self) -> None: ... @property def tags(self) -> Tags: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 329baaa0e..ecf3c417f 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -912,6 +912,10 @@ class RemoteTable(Table): """Not supported on LanceDB Cloud.""" return LOOP.run(self._table.unset_lsm_write_spec()) + def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]: + """Read the installed LsmWriteSpec, or ``None``.""" + return LOOP.run(self._table.get_lsm_write_spec()) + def close_lsm_writers(self) -> None: """No-op on LanceDB Cloud (no local shard writers).""" return LOOP.run(self._table.close_lsm_writers()) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 25df6b554..8ea4f1ea8 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3742,6 +3742,11 @@ class LanceTable(Table): [`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec].""" return LOOP.run(self._table.unset_lsm_write_spec()) + def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]: + """Read the installed LsmWriteSpec, or ``None``. See + [`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec].""" + return LOOP.run(self._table.get_lsm_write_spec()) + def close_lsm_writers(self) -> None: """Close cached MemWAL shard writers. See [`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers].""" @@ -4391,6 +4396,17 @@ class AsyncTable: """ await self._inner.unset_lsm_write_spec() + async def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]: + """Read the LsmWriteSpec currently installed on this table. + + Returns ``None`` when the MemWAL LSM write path is not enabled (no + spec has been set, or it was removed with `unset_lsm_write_spec`). + The returned spec — including its ``maintained_indexes`` and + ``writer_config_defaults`` — mirrors what was passed to + `set_lsm_write_spec`. + """ + return await self._inner.get_lsm_write_spec() + async def close_lsm_writers(self) -> None: """Drain and close any cached MemWAL shard writers for this table. diff --git a/python/python/tests/test_lsm_write_spec.py b/python/python/tests/test_lsm_write_spec.py index d9d75d3a9..9d7acbcb4 100644 --- a/python/python/tests/test_lsm_write_spec.py +++ b/python/python/tests/test_lsm_write_spec.py @@ -136,3 +136,62 @@ def test_lsm_write_spec_identity_and_writer_config_defaults(): s = s.with_writer_config_defaults({"durable_write": "false"}) assert s.writer_config_defaults == {"durable_write": "false"} assert "durable_write" in repr(s) + + +def test_get_lsm_write_spec(tmp_path): + _db, table = _make_table(tmp_path) + table.set_unenforced_primary_key("id") + + # None when nothing is installed. + assert table.get_lsm_write_spec() is None + + # Bucket spec round-trips, including writer config defaults. + table.set_lsm_write_spec( + LsmWriteSpec.bucket("id", 4).with_writer_config_defaults( + {"durable_write": "false"} + ) + ) + spec = table.get_lsm_write_spec() + assert spec is not None + assert spec.spec_type == "bucket" + assert spec.column == "id" + assert spec.num_buckets == 4 + assert spec.writer_config_defaults == {"durable_write": "false"} + + # After unset, None again. + table.unset_lsm_write_spec() + assert table.get_lsm_write_spec() is None + + # Identity round-trips (column recovered from the schema). + table.set_lsm_write_spec(LsmWriteSpec.identity("id")) + spec = table.get_lsm_write_spec() + assert spec.spec_type == "identity" + assert spec.column == "id" + table.unset_lsm_write_spec() + + # Unsharded round-trips (no routing column). + table.set_lsm_write_spec(LsmWriteSpec.unsharded()) + spec = table.get_lsm_write_spec() + assert spec.spec_type == "unsharded" + assert spec.column is None + + +@pytest.mark.asyncio +async def test_async_get_lsm_write_spec(tmp_path): + db = await lancedb.connect_async( + tmp_path, read_consistency_interval=timedelta(seconds=0) + ) + table = await db.create_table( + "t", + pa.RecordBatchReader.from_batches(SCHEMA, [_batch(["seed"], [0])]), + ) + + assert await table.get_lsm_write_spec() is None + await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8)) + spec = await table.get_lsm_write_spec() + assert spec is not None + assert spec.spec_type == "bucket" + assert spec.column == "id" + assert spec.num_buckets == 8 + await table.unset_lsm_write_spec() + assert await table.get_lsm_write_spec() is None diff --git a/python/src/table.rs b/python/src/table.rs index 611f3c9f2..7bc17ccea 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -322,6 +322,12 @@ impl From for lancedb::table::LsmWriteSpec { } } +impl From for LsmWriteSpec { + fn from(inner: lancedb::table::LsmWriteSpec) -> Self { + Self { inner } + } +} + #[pyclass(get_all, from_py_object)] #[derive(Clone, Debug)] pub struct AddColumnsResult { @@ -1029,6 +1035,14 @@ impl Table { }) } + pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let spec = inner.get_lsm_write_spec().await.infer_error()?; + Ok(spec.map(LsmWriteSpec::from)) + }) + } + pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index f3976d4af..b3003a626 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1489,6 +1489,85 @@ impl RemoteTable { } } +/// Reconstruct an [`LsmWriteSpec`](crate::table::LsmWriteSpec) from the MemWAL +/// index's `details` JSON as emitted by the server's `describe_indices`. +/// +/// Mirrors the native reconstruction in `table::merge::lsm`, but reads the +/// server-resolved `column` name from the sharding field rather than a Lance +/// field id — field ids do not travel to the remote client, so the server +/// resolves the name against its schema before serializing the details. +fn lsm_write_spec_from_details_json( + details: &str, +) -> std::result::Result { + use crate::table::LsmWriteSpec; + + #[derive(Deserialize)] + struct Details { + #[serde(default)] + sharding_specs: Vec, + #[serde(default)] + maintained_indexes: Vec, + #[serde(default)] + writer_config_defaults: HashMap, + } + #[derive(Deserialize)] + struct ShardingSpec { + #[serde(default)] + fields: Vec, + } + #[derive(Deserialize)] + struct ShardingField { + transform: Option, + /// Server-resolved routing column name (bucket / identity only). + #[serde(default)] + column: Option, + #[serde(default)] + parameters: HashMap, + } + + let details: Details = serde_json::from_str(details) + .map_err(|e| format!("failed to parse MemWAL index details: {}", e))?; + + let field = details + .sharding_specs + .first() + .and_then(|s| s.fields.first()) + .ok_or_else(|| "MemWAL index details has no sharding field".to_string())?; + + let column = || { + field + .column + .clone() + .ok_or_else(|| "MemWAL sharding field is missing the resolved column name".to_string()) + }; + + let base = match field.transform.as_deref() { + Some("bucket") => { + let num_buckets = field + .parameters + .get("num_buckets") + .and_then(|s| s.parse::().ok()) + .filter(|n| *n > 0) + .ok_or_else(|| { + "MemWAL bucket spec has a missing or invalid num_buckets".to_string() + })?; + LsmWriteSpec::bucket(column()?, num_buckets) + } + Some("identity") => LsmWriteSpec::identity(column()?), + Some("unsharded") => LsmWriteSpec::unsharded(), + other => { + return Err(format!( + "MemWAL index has an unsupported sharding transform {:?}", + other + )); + } + }; + + Ok(base + .with_maintained_indexes(details.maintained_indexes) + .with_writer_config_defaults(details.writer_config_defaults)) +} + #[async_trait] impl BaseTable for RemoteTable { fn as_any(&self) -> &dyn std::any::Any { @@ -2278,6 +2357,65 @@ impl BaseTable for RemoteTable { Ok(()) } + async fn get_lsm_write_spec(&self) -> Result> { + // The spec lives in the `__lance_mem_wal` system index, which the + // curated `index/list` surface filters out by default. `include_system` + // opts that entry back in for this read only — `list_indices` stays + // curated. The system index's `index_details` carries the decoded + // MemWAL state (with the server-resolved shard column name). + let mut request = self.post_read(&format!("/v1/table/{}/index/list/", self.identifier)); + let version = self.current_version().await; + let mut body = serde_json::json!({ "version": version, "include_system": true }); + self.apply_branch_body(&mut body); + request = request.json(&body); + + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + #[derive(Deserialize)] + struct IndexList { + #[serde(default)] + indexes: Vec, + } + // Only the name and opaque details are read; `index_type` is ignored so + // the system index's type (which has no client `IndexType` variant) does + // not break deserialization. + #[derive(Deserialize)] + struct IndexEntry { + index_name: String, + index_details: Option, + } + + let parsed: IndexList = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse index list response: {}", e).into(), + request_id: request_id.clone(), + status_code: None, + })?; + + let Some(entry) = parsed + .indexes + .into_iter() + .find(|i| i.index_name.as_str() == lance_index::mem_wal::MEM_WAL_INDEX_NAME) + else { + // No MemWAL index installed → LSM write path not enabled. + return Ok(None); + }; + + let details = entry.index_details.ok_or_else(|| Error::Http { + source: "index/list returned the MemWAL index without details".into(), + request_id: request_id.clone(), + status_code: None, + })?; + + let spec = lsm_write_spec_from_details_json(&details).map_err(|msg| Error::Http { + source: msg.into(), + request_id, + status_code: None, + })?; + Ok(Some(spec)) + } + async fn tags(&self) -> Result> { Ok(Box::new(RemoteTags { inner: self })) } @@ -5317,6 +5455,82 @@ mod tests { table.unset_lsm_write_spec().await.unwrap(); } + #[tokio::test] + async fn test_get_lsm_write_spec() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/index/list/"); + // The getter opts the system index back into the listing. + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["include_system"], serde_json::json!(true)); + + // The MemWAL index carries the decoded state in `index_details`, + // including the server-resolved `column` name. A user index is + // included too, with a type the client has no variant for — proving + // the getter reads by name and ignores `index_type`. + let details = r#"{"sharding_specs":[{"fields":[{"transform":"bucket","column":"id","parameters":{"num_buckets":"4"}}]}],"maintained_indexes":["id_idx"],"writer_config_defaults":{"durable_write":"false"}}"#; + let response = serde_json::json!({ + "indexes": [ + { "index_name": "v_idx", "index_type": "BTREE", "columns": ["v"] }, + { + "index_name": "__lance_mem_wal", + "index_type": "MemWal", + "columns": [], + "index_details": details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + + let spec = table + .get_lsm_write_spec() + .await + .unwrap() + .expect("a spec should be reported"); + match spec { + crate::table::LsmWriteSpec::Bucket { + column, + num_buckets, + maintained_indexes, + writer_config_defaults, + } => { + assert_eq!(column, "id"); + assert_eq!(num_buckets, 4); + assert_eq!(maintained_indexes, vec!["id_idx".to_string()]); + assert_eq!( + writer_config_defaults + .get("durable_write") + .map(String::as_str), + Some("false") + ); + } + other => panic!("expected a bucket spec, got {:?}", other), + } + } + + #[tokio::test] + async fn test_get_lsm_write_spec_absent() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/index/list/"); + // No __lance_mem_wal entry → the LSM write path is not enabled. + let response = serde_json::json!({ + "indexes": [ + { "index_name": "v_idx", "index_type": "BTREE", "columns": ["v"] }, + ] + }); + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + assert!(table.get_lsm_write_spec().await.unwrap().is_none()); + } + #[tokio::test] async fn test_wait_for_index() { let table = _make_table_with_indices(0); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 156a07821..51d4bd684 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -581,6 +581,16 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "unset_lsm_write_spec is not supported on this table type".into(), }) } + /// Read the [`LsmWriteSpec`] currently installed on this table, returning + /// `None` when the MemWAL LSM write path is not enabled. + /// + /// The default implementation returns `NotSupported`. Implementations that + /// support the MemWAL LSM write path must override this. + async fn get_lsm_write_spec(&self) -> Result> { + Err(Error::NotSupported { + message: "get_lsm_write_spec is not supported on this table type".into(), + }) + } /// Drain and close any cached MemWAL shard writers for this table. /// /// The default implementation is a no-op; table types that maintain @@ -1538,6 +1548,29 @@ impl Table { self.inner.unset_lsm_write_spec().await } + /// Read the [`LsmWriteSpec`] currently installed on this table. + /// + /// Returns `Ok(None)` when the MemWAL LSM write path is not enabled (no + /// spec has been set, or it was removed with [`Table::unset_lsm_write_spec`]). + /// The returned spec — including its [`LsmWriteSpec::maintained_indexes`] and + /// [`LsmWriteSpec::writer_config_defaults`] — mirrors what was passed to + /// [`Table::set_lsm_write_spec`]. + /// + /// # Example + /// + /// ``` + /// # use lancedb::table::Table; + /// # async fn example(table: &Table) -> Result<(), Box> { + /// if let Some(spec) = table.get_lsm_write_spec().await? { + /// println!("LSM write path enabled: {:?}", spec); + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn get_lsm_write_spec(&self) -> Result> { + self.inner.get_lsm_write_spec().await + } + /// Drain and close any cached MemWAL shard writers held for this table. /// /// When an [`LsmWriteSpec`] is installed, `merge_insert` opens MemWAL shard @@ -2850,6 +2883,10 @@ impl BaseTable for NativeTable { merge::lsm::unset_lsm_write_spec(self).await } + async fn get_lsm_write_spec(&self) -> Result> { + merge::lsm::get_lsm_write_spec(self).await + } + async fn close_lsm_writers(&self) -> Result<()> { merge::lsm::close_lsm_writers(self).await } @@ -4411,6 +4448,67 @@ mod tests { } } + #[tokio::test] + async fn test_get_lsm_write_spec() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("region", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone())); + let conn = ConnectBuilder::new(uri) + .read_consistency_interval(Duration::from_secs(0)) + .execute() + .await + .unwrap(); + let table = conn.create_table("t", reader).execute().await.unwrap(); + + // No spec installed yet. + assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); + + // A real scalar index is needed to name it as a maintained index. + table + .create_index(&["id"], Index::Auto) + .execute() + .await + .unwrap(); + let idx_name = table.list_indices().await.unwrap()[0].name.clone(); + + // Bucket spec round-trips exactly, including the routing column (recovered + // from its field id), maintained indexes, and writer config defaults. + let spec = LsmWriteSpec::bucket("id", 4) + .with_maintained_indexes([idx_name]) + .with_writer_config_defaults([("durable_write", "false")]); + table.set_lsm_write_spec(spec.clone()).await.unwrap(); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); + + // After unset, no spec is reported. + table.unset_lsm_write_spec().await.unwrap(); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); + + // Identity sharding round-trips (column recovered from the schema). + let spec = LsmWriteSpec::identity("region"); + table.set_lsm_write_spec(spec.clone()).await.unwrap(); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); + table.unset_lsm_write_spec().await.unwrap(); + + // Unsharded round-trips (no routing column). + let spec = LsmWriteSpec::unsharded(); + table.set_lsm_write_spec(spec.clone()).await.unwrap(); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); + } + #[tokio::test] pub async fn test_stats() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 80246d59f..5577f6128 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -32,7 +32,7 @@ use lance::dataset::mem_wal::{ }; use lance::index::DatasetIndexExt; use lance_core::datatypes::Schema as LanceSchema; -use lance_index::mem_wal::{MemWalIndexDetails, ShardingSpec}; +use lance_index::mem_wal::{MemWalIndexDetails, ShardingField, ShardingSpec}; use tokio::sync::RwLock; use uuid::Uuid; @@ -148,6 +148,97 @@ pub(crate) async fn unset_lsm_write_spec(table: &NativeTable) -> Result<()> { Ok(()) } +// ============================================================================= +// get_lsm_write_spec +// ============================================================================= + +/// Read the [`LsmWriteSpec`] currently installed on the table, if any. +/// +/// Reconstructs the spec from the MemWAL index's stored details — the same +/// metadata [`set_lsm_write_spec`] writes — read directly via +/// [`mem_wal_index_details`](DatasetMemWalExt::mem_wal_index_details) (the raw +/// manifest record, not the `describe_indices` enrichment path, so it is +/// unaffected by system-index filtering there). Returns `Ok(None)` when no +/// MemWAL index is installed. +#[allow(clippy::redundant_pub_crate)] +pub(crate) async fn get_lsm_write_spec(table: &NativeTable) -> Result> { + let dataset = table.dataset.get().await?; + let Some(details) = dataset.mem_wal_index_details().await? else { + return Ok(None); + }; + let spec = lsm_write_spec_from_details(&details, dataset.schema())?; + Ok(Some(spec)) +} + +/// Reconstruct the public [`LsmWriteSpec`] from stored MemWAL index details. +/// +/// The sharding transform and its parameters recover the mode; the routing +/// column is resolved from the sharding field's source id back to its schema +/// name (only bucket / identity carry a column). `maintained_indexes` and +/// `writer_config_defaults` are copied verbatim, so the result round-trips +/// what `set_lsm_write_spec` installed. +fn lsm_write_spec_from_details( + details: &MemWalIndexDetails, + schema: &LanceSchema, +) -> Result { + let spec = details + .sharding_specs + .first() + .ok_or_else(|| Error::Runtime { + message: "get_lsm_write_spec: MemWAL index has no sharding spec".to_string(), + })?; + let field = spec.fields.first().ok_or_else(|| Error::Runtime { + message: "get_lsm_write_spec: MemWAL index has an empty sharding spec".to_string(), + })?; + + let base = match field.transform.as_deref() { + Some(BUCKET_TRANSFORM) => { + let num_buckets = field + .parameters + .get(NUM_BUCKETS_PARAM) + .and_then(|s| s.parse::().ok()) + .filter(|n| *n > 0) + .ok_or_else(|| Error::Runtime { + message: "get_lsm_write_spec: MemWAL bucket spec has a missing or invalid num_buckets parameter".to_string(), + })?; + LsmWriteSpec::bucket(sharding_column(field, schema)?, num_buckets) + } + Some(IDENTITY_TRANSFORM) => LsmWriteSpec::identity(sharding_column(field, schema)?), + Some(UNSHARDED_TRANSFORM) => LsmWriteSpec::unsharded(), + other => { + return Err(Error::Runtime { + message: format!( + "get_lsm_write_spec: MemWAL index has an unsupported sharding transform {:?}", + other + ), + }); + } + }; + + Ok(base + .with_maintained_indexes(details.maintained_indexes.clone()) + .with_writer_config_defaults(details.writer_config_defaults.clone())) +} + +/// Resolve the single routing column name from a sharding field's source id. +/// +/// `set_lsm_write_spec` records the shard column by its Lance field id, so the +/// name is recovered by looking that id up in the current schema. +fn sharding_column(field: &ShardingField, schema: &LanceSchema) -> Result { + let source_id = *field.source_ids.first().ok_or_else(|| Error::Runtime { + message: "get_lsm_write_spec: sharding field has no source column".to_string(), + })?; + schema + .field_by_id(source_id) + .map(|f| f.name.clone()) + .ok_or_else(|| Error::Runtime { + message: format!( + "get_lsm_write_spec: sharding source field id {} not found in schema", + source_id + ), + }) +} + // ============================================================================= // close_lsm_writers // =============================================================================