diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 646fe48c4..4e96fe9d3 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -398,6 +398,26 @@ Drop an index from the table. *** +### getLsmWriteSpec() + +```ts +abstract getLsmWriteSpec(): Promise +``` + +Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) 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 [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)). +The returned spec — including its `maintainedIndexes` and +`writerConfigDefaults` — mirrors what was passed to +[Table#setLsmWriteSpec](Table.md#setlsmwritespec). + +#### Returns + +`Promise`<`undefined` \| [`LsmWriteSpec`](../interfaces/LsmWriteSpec.md)> + +*** + ### indexStats() ```ts diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index f4ba37361..925b0b155 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2992,6 +2992,56 @@ 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(); + + // A real scalar index is needed to name it as a maintained index. + await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]); + await table.createIndex("id"); + const indexName = (await table.listIndices())[0].name; + + // Bucket spec round-trips, including maintained indexes and writer config + // defaults. Lance writer-config keys are canonically snake_case. + // biome-ignore lint/style/useNamingConvention: Lance writer-config keys are snake_case + const writerConfigDefaults = { durable_write: "false" }; + await table.setLsmWriteSpec({ + specType: "bucket", + column: "id", + numBuckets: 4, + maintainedIndexes: [indexName], + writerConfigDefaults, + }); + 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?.maintainedIndexes).toEqual([indexName]); + expect(spec?.writerConfigDefaults).toEqual(writerConfigDefaults); + + // 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..84318e63c 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,15 @@ export class LocalTable extends Table { return await this.inner.unsetLsmWriteSpec(); } + async getLsmWriteSpec(): Promise { + // The native binding types `specType` as a plain `string`; narrow it back + // to the public union. The Rust `From` impl only ever emits one of the + // three valid values, so the cast is safe. + return ((await this.inner.getLsmWriteSpec()) ?? undefined) as + | LsmWriteSpec + | 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 e71169624..8667b32c3 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -229,6 +229,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 f4d70077c..bc11a4ce7 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3759,6 +3759,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].""" @@ -4417,6 +4422,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..d38918f09 100644 --- a/python/python/tests/test_lsm_write_spec.py +++ b/python/python/tests/test_lsm_write_spec.py @@ -11,6 +11,7 @@ import lancedb import pyarrow as pa import pytest from lancedb._lancedb import LsmWriteSpec +from lancedb.index import BTree SCHEMA = pa.schema( [ @@ -136,3 +137,76 @@ 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 + + # A real scalar index is needed to name it as a maintained index. + table.create_index("id", config=BTree()) + idx_name = table.list_indices()[0].name + + # Bucket spec round-trips, including maintained indexes and writer config + # defaults. + table.set_lsm_write_spec( + LsmWriteSpec.bucket("id", 4) + .with_maintained_indexes([idx_name]) + .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.maintained_indexes == [idx_name] + 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 + + # A real scalar index is needed to name it as a maintained index. + await table.create_index("id", config=BTree()) + idx_name = (await table.list_indices())[0].name + + await table.set_lsm_write_spec( + LsmWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name]) + ) + 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 + assert spec.maintained_indexes == [idx_name] + 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 29495b73d..c9ec28ee4 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -20,6 +20,7 @@ use crate::table::AddColumnsResult; use crate::table::AddResult; use crate::table::DeleteResult; use crate::table::DropColumnsResult; +use crate::table::LsmWriteSpec; use crate::table::MergeResult; use crate::table::Tags; use crate::table::UpdateResult; @@ -2269,8 +2270,7 @@ impl BaseTable for RemoteTable { }) } - async fn set_lsm_write_spec(&self, spec: crate::table::LsmWriteSpec) -> Result<()> { - use crate::table::LsmWriteSpec; + async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> { self.check_mutable().await?; // Map the spec onto the server's request DTO. `sharding` is internally @@ -2322,6 +2322,69 @@ impl BaseTable for RemoteTable { Ok(()) } + async fn get_lsm_write_spec(&self) -> Result> { + // Read counterpart to set/unset, resolved server-side against HEAD. The + // server reads the spec from the `__lance_mem_wal` system index (shard + // column mapped from its Lance field id against the current schema) and + // re-encodes it into the same sophon-owned shape the set endpoint + // accepts — no lance/lancedb types cross the wire. `lsm_write_spec` is + // null when the LSM write path is not enabled for the table. + let request = self.post_read(&format!( + "/v1/table/{}/get_lsm_write_spec/", + self.identifier + )); + 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())?; + + // Mirror of sophon's `Sharding` (internally tagged on `mode`) and + // `LsmWriteSpecBody` / `GetLsmWriteSpecResponse`. + #[derive(Deserialize)] + #[serde(tag = "mode", rename_all = "snake_case")] + enum Sharding { + Unsharded, + Bucket { column: String, num_buckets: u32 }, + Identity { column: String }, + } + #[derive(Deserialize)] + struct LsmWriteSpecBody { + sharding: Sharding, + #[serde(default)] + maintained_indexes: Vec, + #[serde(default)] + writer_config_defaults: std::collections::HashMap, + } + #[derive(Deserialize)] + struct GetLsmWriteSpecResponse { + lsm_write_spec: Option, + } + + let parsed: GetLsmWriteSpecResponse = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse get_lsm_write_spec response: {}", e).into(), + request_id, + status_code: None, + })?; + + let Some(body) = parsed.lsm_write_spec else { + // The LSM write path is not enabled for this table. + return Ok(None); + }; + + let spec = match body.sharding { + Sharding::Bucket { + column, + num_buckets, + } => LsmWriteSpec::bucket(column, num_buckets), + Sharding::Identity { column } => LsmWriteSpec::identity(column), + Sharding::Unsharded => LsmWriteSpec::unsharded(), + } + .with_maintained_indexes(body.maintained_indexes) + .with_writer_config_defaults(body.writer_config_defaults); + + Ok(Some(spec)) + } + async fn tags(&self) -> Result> { Ok(Box::new(RemoteTags { inner: self })) } @@ -5361,6 +5424,74 @@ 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/get_lsm_write_spec/" + ); + + // The server resolves the spec and re-encodes it into the same + // sophon-owned shape the set endpoint accepts (`Sharding` internally + // tagged on `mode`, wrapped in `lsm_write_spec`). + let response = serde_json::json!({ + "lsm_write_spec": { + "sharding": { "mode": "bucket", "column": "id", "num_buckets": 4 }, + "maintained_indexes": ["id_idx"], + "writer_config_defaults": { "durable_write": "false" }, + } + }); + 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/get_lsm_write_spec/" + ); + // Null spec → the LSM write path is not enabled. + let response = serde_json::json!({ "lsm_write_spec": null }); + 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 // =============================================================================