mirror of
https://github.com/lancedb/lancedb.git
synced 2026-07-10 06:20:48 +00:00
feat: add get_lsm_write_spec to read the installed LSM write spec (#3631)
## Summary Adds `Table::get_lsm_write_spec` returning `Option<LsmWriteSpec>` — the read counterpart to the existing `set_lsm_write_spec` / `unset_lsm_write_spec`. Returns `None` when the MemWAL LSM write path is not enabled; otherwise reconstructs the spec (mode, shard column, `num_buckets`, `maintained_indexes`, `writer_config_defaults`) exactly as installed. ## Changes - **Rust core (`NativeTable`)** — reconstructs the spec from `mem_wal_index_details()`, resolving the shard column from its Lance field id via the dataset schema. This is a raw metadata read, so it is unaffected by `describe_indices` system-index filtering. - **Remote (`RemoteTable`)** — reads the `__lance_mem_wal` system index through `index/list` with `include_system: true` (so the curated `list_indices` surface stays unchanged), then parses the index `details` JSON. It matches the index by name and ignores `index_type`, so no client `IndexType` variant is needed. It uses the **server-resolved `column` name** from the details (Lance field ids do not travel to the remote client). - **Python + TypeScript bindings** — sync and async, mirroring `set`/`unset`, with round-trip tests (bucket / identity / unsharded, plus `None` when unset). ## Tests - Rust: native round-trip unit test + remote mock-endpoint tests (present + absent). All green (`cargo test --features remote -p lancedb`). - Python/TS: round-trip tests added; binding-runtime execution runs in CI. ## Dependencies for the remote path The remote path is complete on the client side but depends on two out-of-repo pieces to work end-to-end: 1. **lance** — emit the server-resolved shard **`column`** name in the MemWAL index `details` JSON (field ids can't reach the client). See lance-format/lance#7667. 2. **server** — honor `include_system` on `index/list` so the `__lance_mem_wal` entry is returned for this read. Against an older server (no `include_system`), the remote getter degrades gracefully to `Ok(None)` rather than erroring. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -398,6 +398,26 @@ Drop an index from the table.
|
||||
|
||||
***
|
||||
|
||||
### getLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
abstract getLsmWriteSpec(): Promise<undefined | LsmWriteSpec>
|
||||
```
|
||||
|
||||
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
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -585,6 +585,17 @@ export abstract class Table {
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
abstract unsetLsmWriteSpec(): Promise<void>;
|
||||
/**
|
||||
* 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<LsmWriteSpec | undefined>}
|
||||
*/
|
||||
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
|
||||
/**
|
||||
* 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<LsmWriteSpec | undefined> {
|
||||
// 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<void> {
|
||||
return await this.inner.closeLsmWriters();
|
||||
}
|
||||
|
||||
@@ -411,6 +411,16 @@ impl Table {
|
||||
.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn get_lsm_write_spec(&self) -> napi::Result<Option<LsmWriteSpec>> {
|
||||
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<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> 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)]
|
||||
|
||||
@@ -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: ...
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -322,6 +322,12 @@ impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> 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<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
|
||||
@@ -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<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
})
|
||||
}
|
||||
|
||||
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<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
|
||||
// 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<String>,
|
||||
#[serde(default)]
|
||||
writer_config_defaults: std::collections::HashMap<String, String>,
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct GetLsmWriteSpecResponse {
|
||||
lsm_write_spec: Option<LsmWriteSpecBody>,
|
||||
}
|
||||
|
||||
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<Box<dyn Tags + '_>> {
|
||||
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);
|
||||
|
||||
@@ -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<Option<LsmWriteSpec>> {
|
||||
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<dyn std::error::Error>> {
|
||||
/// 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<Option<LsmWriteSpec>> {
|
||||
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<Option<LsmWriteSpec>> {
|
||||
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<dyn arrow_array::RecordBatchReader + Send> =
|
||||
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();
|
||||
|
||||
@@ -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<Option<LsmWriteSpec>> {
|
||||
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<LsmWriteSpec> {
|
||||
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::<u32>().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<String> {
|
||||
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
|
||||
// =============================================================================
|
||||
|
||||
Reference in New Issue
Block a user