mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 15:22:30 +00:00
feat: maintain memtable index kinds from a table's registry (#4477)
## Why lance-format/lance#9803 makes the MemWAL's memtable index kinds pluggable: a `MemIndexRegistry` on the writer config decides which base-table index kinds are maintained in memory, so a kind Lance does not build in can be added by registering a plugin. LanceDB opens the MemWAL writer itself for `merge_insert`, and checks a named maintained set itself, both against Lance's built-ins only. So a kind registered outside Lance could not be named in `LsmWriteSpec::with_maintained_indexes`, and under "maintain every index" it was skipped. ## What changes - `NativeTable::set_mem_index_registry(registry)` sets the memtable index kinds a table maintains. The default is Lance's built-ins, so nothing changes unless a caller sets one. The registry checks a named maintained set (`validate_maintained_indexes_with`), is recorded on `initialize_mem_wal`, and goes on the writer config `merge_insert` opens. `lancedb::table` re-exports `MemIndexRegistry` and Lance's plugin traits (`mem_index`), so a caller needs no direct `lance` dependency. - The LSM vector read no longer checks resident memtables at all, and the `carries_vector_index` helper goes with it. Lance already compares every row of a memtable whose index can't answer the search, so no rows are missed. It still requires the base vector index to be maintained. ## Limits - The registry belongs to the table handle. Clones of a handle and its branch handles share it; a handle opened separately starts from the built-ins. Set it before the handle's first LSM write: a writer already open keeps the kinds it opened with. - The maintained set is stored with the table but the registry is not. Once a spec names an index only a registered plugin maintains, an LSM write from a handle without that registry is refused with Lance's "no registered plugin maintains" error rather than skipping the index. That includes the Python and Node bindings, which have no way to set a registry yet. - Rust only. A registry holds Rust plugin objects, which the Python, TypeScript and Java bindings cannot construct, so they get no setter; their documented behaviour (built-in kinds only) is unchanged. - Per-index plugin settings (`ShardWriterConfig::index_overrides`, from lance-format/lance#9805) are not exposed; writers use each plugin's defaults. ## Open decision Whether the registry should also be settable on the connection, so every table a connection opens carries it. This PR keeps it per table; a connection-level option would sit on top of `set_mem_index_registry` without changing it. ## Depends on Lance v14.0.0-beta.13, which `main` already uses. It includes lance-format/lance#9803 and #9805 (pluggable memtable index kinds) and #9857 (flushed generations keep the table's field ids). ## Testing - `lsm_maintains_an_index_kind_from_the_tables_registry`: a stand-in plugin maintains the base table's bitmap index as a B-tree. Naming the index is refused without the registry and accepted with it; a separately opened handle without the registry is refused an LSM write. Rows written through `merge_insert` are found by id before and after the flush. Each generation the shard manifests publish carries a bitmap index on `region`, and a filter on it is answered through that index. - `lsm_default_set_maintains_a_registry_kind_created_while_writing`: under "maintain every index", a bitmap index created after the writer opened is picked up by that writer and flushed. A clone of the table handle shares the registry. - `lsm_vector_search_compares_every_row_when_no_resident_index_answers`: with a vector plugin that answers no search, LSM vector search still returns the right rows, including a second page. - `cargo test -p lancedb --features remote --lib` passes; `cargo clippy -p lancedb --features remote --tests -- -D warnings` and `cargo fmt --all -- --check` are clean. --------- Co-authored-by: XYZhan <xuanyu.zhan@lancedb.com>
This commit is contained in:
@@ -19,6 +19,9 @@ pub use lance::dataset::ReadParams;
|
||||
pub use lance::dataset::Version;
|
||||
use lance::dataset::WriteMode;
|
||||
use lance::dataset::builder::DatasetBuilder;
|
||||
pub use lance::dataset::mem_wal::MemIndexRegistry;
|
||||
/// The traits a memtable index plugin implements; see [`MemIndexRegistry`].
|
||||
pub use lance::dataset::mem_wal::index as mem_index;
|
||||
use lance::dataset::{InsertBuilder, WriteParams};
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance::index::scalar::load_segment_params;
|
||||
@@ -424,7 +427,8 @@ impl LsmWriteSpec {
|
||||
/// changes, so an index created later is maintained too; one of a kind the
|
||||
/// MemWAL cannot mirror is skipped rather than failing the table. A list is
|
||||
/// verbatim: each name must already exist and be maintainable, and an empty
|
||||
/// list maintains nothing.
|
||||
/// list maintains nothing. The maintainable kinds are Lance's built-ins plus
|
||||
/// any added with [`NativeTable::set_mem_index_registry`].
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::table::LsmWriteSpec;
|
||||
@@ -2590,6 +2594,9 @@ pub struct NativeTable {
|
||||
pub(crate) pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
|
||||
// Read-freshness baseline; `Some` only for namespace-backed tables.
|
||||
freshness: Option<TableFreshness>,
|
||||
// Memtable index kinds this handle's LSM writes maintain, shared with its
|
||||
// clones and branch handles.
|
||||
mem_index_registry: Arc<std::sync::RwLock<MemIndexRegistry>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for NativeTable {
|
||||
@@ -2756,9 +2763,51 @@ impl NativeTable {
|
||||
namespace_client,
|
||||
pushdown_operations,
|
||||
freshness: None,
|
||||
mem_index_registry: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Maintain the memtable index kinds in `registry` when this table writes
|
||||
/// through its [`LsmWriteSpec`], and accept them when
|
||||
/// [`LsmWriteSpec::with_maintained_indexes`] names one.
|
||||
///
|
||||
/// The registry belongs to this handle, its clones and its branch handles,
|
||||
/// and is not stored with the table: a handle opened separately starts from
|
||||
/// Lance's built-ins. Set it before the handle's first LSM write; writers
|
||||
/// already open keep the registry they were opened with. Local tables only,
|
||||
/// reached through [`Table::as_native`].
|
||||
///
|
||||
/// ```
|
||||
/// # use std::sync::Arc;
|
||||
/// # use lancedb::Table;
|
||||
/// # use lancedb::table::{LsmWriteSpec, MemIndexRegistry, mem_index::MemIndexPlugin};
|
||||
/// # async fn example(
|
||||
/// # table: &Table,
|
||||
/// # plugin: Arc<dyn MemIndexPlugin>,
|
||||
/// # ) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let registry = MemIndexRegistry::default().with_plugin(plugin)?;
|
||||
/// if let Some(native) = table.as_native() {
|
||||
/// native.set_mem_index_registry(registry);
|
||||
/// }
|
||||
/// let spec = LsmWriteSpec::unsharded().with_maintained_indexes(vec!["region_bitmap".to_string()]);
|
||||
/// table.set_lsm_write_spec(spec).await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn set_mem_index_registry(&self, registry: MemIndexRegistry) {
|
||||
*self
|
||||
.mem_index_registry
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = registry;
|
||||
}
|
||||
|
||||
pub(crate) fn mem_index_registry(&self) -> MemIndexRegistry {
|
||||
self.mem_index_registry
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Set the namespace client for server-side query execution.
|
||||
///
|
||||
/// When set, queries will be executed on the namespace server instead of locally.
|
||||
@@ -2786,6 +2835,7 @@ impl NativeTable {
|
||||
namespace_client: self.namespace_client.clone(),
|
||||
pushdown_operations: self.pushdown_operations.clone(),
|
||||
freshness: self.freshness.clone(),
|
||||
mem_index_registry: self.mem_index_registry.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2893,6 +2943,7 @@ impl NativeTable {
|
||||
namespace_client: stored_namespace_client,
|
||||
pushdown_operations,
|
||||
freshness: None,
|
||||
mem_index_registry: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2984,6 +3035,7 @@ impl NativeTable {
|
||||
namespace_client,
|
||||
pushdown_operations,
|
||||
freshness: None,
|
||||
mem_index_registry: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -3117,6 +3169,7 @@ impl NativeTable {
|
||||
namespace_client: stored_namespace_client,
|
||||
pushdown_operations,
|
||||
freshness: None,
|
||||
mem_index_registry: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+437
-21
@@ -1782,8 +1782,13 @@ mod lsm_tests {
|
||||
assert_eq!(rows, 1, "LSM FTS must still see base rows");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lsm_read_vector_search() {
|
||||
/// A table whose vector index the LSM maintains: base rows filled with
|
||||
/// their id, and one memtable row filled with 1000, nearest to `[1000; 8]`.
|
||||
/// `registry` replaces the table's memtable kinds.
|
||||
async fn lsm_vector_table(
|
||||
dir: &std::path::Path,
|
||||
registry: Option<lance::dataset::mem_wal::MemIndexRegistry>,
|
||||
) -> Table {
|
||||
use crate::index::Index;
|
||||
use crate::index::vector::IvfPqIndexBuilder;
|
||||
use arrow::array::{FixedSizeListBuilder, Float32Builder};
|
||||
@@ -1815,12 +1820,7 @@ mod lsm_tests {
|
||||
.unwrap()
|
||||
};
|
||||
|
||||
let dir = tempdir().unwrap();
|
||||
let conn = connect(dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
// Base rows fill each vector with its own id (0..256); all far from 1000.
|
||||
let conn = connect(dir.to_str().unwrap()).execute().await.unwrap();
|
||||
let base = make_batch((0..N).map(|i| (i, i as f32)).collect());
|
||||
let base_reader: Box<dyn RecordBatchReader + Send> =
|
||||
Box::new(RecordBatchIterator::new(vec![Ok(base)], schema.clone()));
|
||||
@@ -1838,13 +1838,15 @@ mod lsm_tests {
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
if let Some(registry) = registry {
|
||||
table.as_native().unwrap().set_mem_index_registry(registry);
|
||||
}
|
||||
let vec_index = table.list_indices().await.unwrap()[0].name.clone();
|
||||
table
|
||||
.set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(vec![vec_index]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Insert a vector (filled with 1000) that is nearest to the query.
|
||||
let mut builder = table.merge_insert(&[]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
@@ -1854,20 +1856,21 @@ mod lsm_tests {
|
||||
schema.clone(),
|
||||
));
|
||||
builder.execute(insert_reader).await.unwrap();
|
||||
table
|
||||
}
|
||||
|
||||
// KNN near [1000; DIM]: the default (auto-routed) read surfaces the
|
||||
// memtable row.
|
||||
let stream = table
|
||||
.query()
|
||||
.nearest_to(&[1000.0_f32; 8])
|
||||
.unwrap()
|
||||
.limit(1)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let ids = collect_ids(stream).await;
|
||||
/// The ids of the nearest row to `[1000; 8]` through the default read.
|
||||
async fn nearest_ids(table: &Table) -> Vec<i64> {
|
||||
let query = table.query().nearest_to(&[1000.0_f32; 8]).unwrap().limit(1);
|
||||
collect_ids(query.execute().await.unwrap()).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lsm_read_vector_search() {
|
||||
let dir = tempdir().unwrap();
|
||||
let table = lsm_vector_table(dir.path(), None).await;
|
||||
assert_eq!(
|
||||
ids,
|
||||
nearest_ids(&table).await,
|
||||
vec![9999],
|
||||
"LSM vector search must rank the memtable row first"
|
||||
);
|
||||
@@ -1893,4 +1896,417 @@ mod lsm_tests {
|
||||
assert_eq!(ids, vec![255]);
|
||||
}
|
||||
}
|
||||
|
||||
/// The built-in vector plugin, except that its index answers no search.
|
||||
#[derive(Debug)]
|
||||
struct AnswersNothing;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct AnswersNothingIndex(Arc<dyn lance::dataset::mem_wal::index::MemIndex>);
|
||||
|
||||
fn vector_plugin() -> Arc<dyn lance::dataset::mem_wal::index::MemIndexPlugin> {
|
||||
Arc::new(lance::dataset::mem_wal::index::HnswMemIndexPlugin)
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl lance::dataset::mem_wal::index::MemIndexPlugin for AnswersNothing {
|
||||
fn name(&self) -> &str {
|
||||
"AnswersNothing"
|
||||
}
|
||||
fn details_message(&self) -> &str {
|
||||
"VectorIndexDetails"
|
||||
}
|
||||
fn flush_index_type(&self) -> lance_index::IndexType {
|
||||
vector_plugin().flush_index_type()
|
||||
}
|
||||
fn training_criteria(&self) -> lance_index::scalar::registry::TrainingCriteria {
|
||||
vector_plugin().training_criteria()
|
||||
}
|
||||
async fn resolve(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::ResolveContext<'_>,
|
||||
) -> lance_core::Result<lance::dataset::mem_wal::index::ResolvedIndex> {
|
||||
vector_plugin().resolve(ctx).await
|
||||
}
|
||||
fn validate(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>,
|
||||
) -> lance_core::Result<()> {
|
||||
vector_plugin().validate(ctx)
|
||||
}
|
||||
fn create(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>,
|
||||
) -> lance_core::Result<Arc<dyn lance::dataset::mem_wal::index::MemIndex>> {
|
||||
Ok(Arc::new(AnswersNothingIndex(vector_plugin().create(ctx)?)))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl lance::dataset::mem_wal::index::MemIndex for AnswersNothingIndex {
|
||||
fn columns(&self) -> &[String] {
|
||||
self.0.columns()
|
||||
}
|
||||
fn can_answer(&self, _query: &dyn lance::dataset::mem_wal::index::MemQuery) -> bool {
|
||||
false
|
||||
}
|
||||
fn insert(&self, batch: &RecordBatch, row_offset: u64) -> lance_core::Result<()> {
|
||||
self.0.insert(batch, row_offset)
|
||||
}
|
||||
fn insert_batches(
|
||||
&self,
|
||||
batches: &[lance::dataset::mem_wal::write::StoredBatch],
|
||||
) -> lance_core::Result<()> {
|
||||
self.0.insert_batches(batches)
|
||||
}
|
||||
fn resident_bytes(&self) -> usize {
|
||||
self.0.resident_bytes()
|
||||
}
|
||||
fn search(
|
||||
&self,
|
||||
_query: &dyn lance::dataset::mem_wal::index::MemQuery,
|
||||
_ctx: &lance::dataset::mem_wal::index::SearchContext,
|
||||
) -> lance_core::Result<Option<lance::dataset::mem_wal::index::MemMatches>> {
|
||||
Ok(None)
|
||||
}
|
||||
async fn flush(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::FlushContext<'_>,
|
||||
) -> lance_core::Result<lance::dataset::mem_wal::index::FlushOutcome> {
|
||||
self.0.flush(ctx).await
|
||||
}
|
||||
}
|
||||
|
||||
/// A resident memtable whose vector index answers nothing still serves an
|
||||
/// LSM vector search: Lance compares every row instead.
|
||||
#[tokio::test]
|
||||
async fn lsm_vector_search_compares_every_row_when_no_resident_index_answers() {
|
||||
let dir = tempdir().unwrap();
|
||||
let mut registry = lance::dataset::mem_wal::MemIndexRegistry::default();
|
||||
registry.replace_plugin(Arc::new(AnswersNothing)).unwrap();
|
||||
let table = lsm_vector_table(dir.path(), Some(registry)).await;
|
||||
assert_eq!(nearest_ids(&table).await, vec![9999]);
|
||||
|
||||
// The second page starts after the memtable row.
|
||||
let page = table
|
||||
.query()
|
||||
.nearest_to(&[1000.0_f32; 8])
|
||||
.unwrap()
|
||||
.limit(2)
|
||||
.offset(1);
|
||||
assert_eq!(
|
||||
collect_ids(page.execute().await.unwrap()).await,
|
||||
vec![255, 254]
|
||||
);
|
||||
}
|
||||
|
||||
/// Claims the base table's bitmap index and keeps it as a B-tree, whose
|
||||
/// flush rows are in the shape the bitmap trainer takes.
|
||||
#[derive(Debug)]
|
||||
struct BitmapAsBTree;
|
||||
|
||||
fn btree_plugin() -> Arc<dyn lance::dataset::mem_wal::index::MemIndexPlugin> {
|
||||
lance::dataset::mem_wal::MemIndexRegistry::default()
|
||||
.plugin_for_details_url("/lance.table.BTreeIndexDetails")
|
||||
.unwrap()
|
||||
.clone()
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl lance::dataset::mem_wal::index::MemIndexPlugin for BitmapAsBTree {
|
||||
fn name(&self) -> &str {
|
||||
"BitmapAsBTree"
|
||||
}
|
||||
fn details_message(&self) -> &str {
|
||||
"BitmapIndexDetails"
|
||||
}
|
||||
fn flush_index_type(&self) -> lance_index::IndexType {
|
||||
lance_index::IndexType::Bitmap
|
||||
}
|
||||
fn training_criteria(&self) -> lance_index::scalar::registry::TrainingCriteria {
|
||||
btree_plugin().training_criteria()
|
||||
}
|
||||
fn validate(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>,
|
||||
) -> lance_core::Result<()> {
|
||||
btree_plugin().validate(ctx)
|
||||
}
|
||||
fn create(
|
||||
&self,
|
||||
ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>,
|
||||
) -> lance_core::Result<Arc<dyn lance::dataset::mem_wal::index::MemIndex>> {
|
||||
btree_plugin().create(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
/// A kind Lance has no memtable index for is refused when named until the
|
||||
/// table's registry maintains it; then it is maintained, flushed with each
|
||||
/// generation, and answers searches on the flushed generations.
|
||||
#[tokio::test]
|
||||
async fn lsm_maintains_an_index_kind_from_the_tables_registry() {
|
||||
let dir = tempdir().unwrap();
|
||||
let conn = connect(dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let table = conn
|
||||
.create_table("t", id_region_reader(vec![(1, "us"), (2, "eu"), (3, "us")]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
table.set_unenforced_primary_key(["id"]).await.unwrap();
|
||||
table
|
||||
.create_index(&["region"], crate::index::Index::Bitmap(Default::default()))
|
||||
.name("region_bitmap".to_string())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let spec =
|
||||
|| LsmWriteSpec::unsharded().with_maintained_indexes(vec!["region_bitmap".to_string()]);
|
||||
|
||||
let refused = table.set_lsm_write_spec(spec()).await.unwrap_err();
|
||||
assert!(
|
||||
refused.to_string().contains("region_bitmap")
|
||||
&& refused
|
||||
.to_string()
|
||||
.contains("no registered plugin maintains"),
|
||||
"{refused:?}"
|
||||
);
|
||||
|
||||
let registry = lance::dataset::mem_wal::MemIndexRegistry::default()
|
||||
.with_plugin(Arc::new(BitmapAsBTree))
|
||||
.unwrap();
|
||||
table.as_native().unwrap().set_mem_index_registry(registry);
|
||||
table.set_lsm_write_spec(spec()).await.unwrap();
|
||||
|
||||
// The registry is not stored with the table.
|
||||
let other = conn.open_table("t").execute().await.unwrap();
|
||||
let mut builder = other.merge_insert(&["id"]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
let refused = builder
|
||||
.execute(id_region_reader(vec![(20, "us")]))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
refused
|
||||
.to_string()
|
||||
.contains("no registered plugin maintains"),
|
||||
"{refused:?}"
|
||||
);
|
||||
|
||||
let mut builder = table.merge_insert(&["id"]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
builder
|
||||
.execute(id_region_reader(vec![
|
||||
(10, "fresh"),
|
||||
(11, "fresh"),
|
||||
(12, "eu"),
|
||||
]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let fresh = || async {
|
||||
let query = table.query().only_if("region = 'fresh'");
|
||||
let mut ids = collect_ids(query.execute().await.unwrap()).await;
|
||||
ids.sort_unstable();
|
||||
ids
|
||||
};
|
||||
assert_eq!(fresh().await, vec![10, 11]);
|
||||
let (_, _, memtables) = table
|
||||
.as_native()
|
||||
.unwrap()
|
||||
.dataset
|
||||
.shard_writer()
|
||||
.read_snapshot()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("the merge opened a writer");
|
||||
let active = memtables.expect("the writer keeps memtables").active;
|
||||
assert!(
|
||||
active
|
||||
.index_store
|
||||
.index_names()
|
||||
.contains(&"region_bitmap".to_string()),
|
||||
"the memtable maintains the plugin's index: {:?}",
|
||||
active.index_store.index_names()
|
||||
);
|
||||
|
||||
table.close_lsm_writers().await.unwrap();
|
||||
assert_eq!(fresh().await, vec![10, 11]);
|
||||
assert_generations_answer_through_region_bitmap(&table, vec![10, 11]).await;
|
||||
}
|
||||
|
||||
/// The generations each shard's manifest publishes.
|
||||
async fn published_generations(table: &Table) -> Vec<lance::Dataset> {
|
||||
use lance::dataset::mem_wal::ShardManifestStore;
|
||||
|
||||
let uri = table.uri().await.unwrap();
|
||||
let (store, base) = lance_io::object_store::ObjectStore::from_uri(&uri)
|
||||
.await
|
||||
.unwrap();
|
||||
let mem_wal = std::path::Path::new(&uri).join("_mem_wal");
|
||||
let mut generations = Vec::new();
|
||||
for shard in std::fs::read_dir(&mem_wal).unwrap() {
|
||||
let Ok(shard_id) = shard.unwrap().file_name().to_string_lossy().parse() else {
|
||||
continue;
|
||||
};
|
||||
let manifest = ShardManifestStore::new(store.clone(), &base, shard_id, 2)
|
||||
.latest()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("a shard directory has a manifest");
|
||||
for sstable in manifest.sstables {
|
||||
let path = mem_wal.join(shard_id.to_string()).join(&sstable.path);
|
||||
generations.push(lance::Dataset::open(path.to_str().unwrap()).await.unwrap());
|
||||
}
|
||||
}
|
||||
generations
|
||||
}
|
||||
|
||||
/// Every published generation carries `region_bitmap` as a bitmap index on
|
||||
/// `region`, and a search for `region = 'fresh'` on it uses that index.
|
||||
async fn assert_generations_answer_through_region_bitmap(table: &Table, fresh: Vec<i64>) {
|
||||
use lance::index::DatasetIndexExt;
|
||||
|
||||
let generations = published_generations(table).await;
|
||||
assert!(
|
||||
!generations.is_empty(),
|
||||
"closing the writer flushed a generation"
|
||||
);
|
||||
let mut found = Vec::new();
|
||||
for generation in generations {
|
||||
let region = generation.schema().field("region").unwrap().id;
|
||||
let indices = generation.load_indices().await.unwrap();
|
||||
let index = indices
|
||||
.iter()
|
||||
.find(|index| index.name == "region_bitmap")
|
||||
.expect("the generation carries the index the plugin trained");
|
||||
assert_eq!(index.fields, vec![region]);
|
||||
assert!(
|
||||
index
|
||||
.index_details
|
||||
.as_ref()
|
||||
.is_some_and(|details| details.type_url.ends_with("BitmapIndexDetails")),
|
||||
"{:?}",
|
||||
index.index_details
|
||||
);
|
||||
|
||||
let mut scan = generation.scan();
|
||||
scan.filter("region = 'fresh'").unwrap();
|
||||
let plan = scan.explain_plan(false).await.unwrap();
|
||||
assert!(plan.contains("ScalarIndexQuery"), "{plan}");
|
||||
let batch = scan.try_into_batch().await.unwrap();
|
||||
found.extend(
|
||||
batch["id"]
|
||||
.as_primitive::<Int64Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.copied(),
|
||||
);
|
||||
}
|
||||
found.sort_unstable();
|
||||
assert_eq!(found, fresh);
|
||||
}
|
||||
|
||||
/// Under the default set, a kind only the registry maintains is picked up
|
||||
/// when its index is created after the writer opened, and the table's
|
||||
/// clones share the registry.
|
||||
#[tokio::test]
|
||||
async fn lsm_default_set_maintains_a_registry_kind_created_while_writing() {
|
||||
let dir = tempdir().unwrap();
|
||||
let conn = connect(dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let table = conn
|
||||
.create_table("t", id_region_reader(vec![(1, "us"), (2, "eu")]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
table.set_unenforced_primary_key(["id"]).await.unwrap();
|
||||
let registry = lance::dataset::mem_wal::MemIndexRegistry::default()
|
||||
.with_plugin(Arc::new(BitmapAsBTree))
|
||||
.unwrap();
|
||||
table.as_native().unwrap().set_mem_index_registry(registry);
|
||||
let clone = table.clone();
|
||||
assert!(
|
||||
clone
|
||||
.as_native()
|
||||
.unwrap()
|
||||
.mem_index_registry()
|
||||
.plugin_for_details_url("/lance.table.BitmapIndexDetails")
|
||||
.is_some(),
|
||||
"a clone shares the handle's registry"
|
||||
);
|
||||
table
|
||||
.set_lsm_write_spec(LsmWriteSpec::unsharded())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let upsert = |rows: Vec<(i64, &'static str)>| {
|
||||
let table = table.clone();
|
||||
async move {
|
||||
let mut builder = table.merge_insert(&["id"]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
builder.execute(id_region_reader(rows)).await.unwrap();
|
||||
}
|
||||
};
|
||||
// Opens the writer before the table has a bitmap index.
|
||||
upsert(vec![(3, "us")]).await;
|
||||
table
|
||||
.create_index(&["region"], crate::index::Index::Bitmap(Default::default()))
|
||||
.name("region_bitmap".to_string())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
upsert(vec![(10, "fresh"), (11, "fresh")]).await;
|
||||
|
||||
let (_, _, memtables) = table
|
||||
.as_native()
|
||||
.unwrap()
|
||||
.dataset
|
||||
.shard_writer()
|
||||
.read_snapshot()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("the merge opened a writer");
|
||||
let active = memtables.expect("the writer keeps memtables").active;
|
||||
assert!(
|
||||
active
|
||||
.index_store
|
||||
.index_names()
|
||||
.contains(&"region_bitmap".to_string()),
|
||||
"the open writer picked up the new index: {:?}",
|
||||
active.index_store.index_names()
|
||||
);
|
||||
|
||||
table.close_lsm_writers().await.unwrap();
|
||||
// Only the generation written after the index exists carries it.
|
||||
let generations = published_generations(&table).await;
|
||||
let mut carrying = 0;
|
||||
for generation in &generations {
|
||||
use lance::index::DatasetIndexExt;
|
||||
if generation
|
||||
.load_indices()
|
||||
.await
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|index| index.name == "region_bitmap")
|
||||
{
|
||||
carrying += 1;
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
carrying >= 1,
|
||||
"{} generations, none carry the index",
|
||||
generations.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,7 +29,7 @@ use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef};
|
||||
use lance::Dataset;
|
||||
use lance::dataset::mem_wal::{
|
||||
DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec,
|
||||
validate_maintained_indexes,
|
||||
validate_maintained_indexes_with,
|
||||
};
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance_core::datatypes::Schema as LanceSchema;
|
||||
@@ -103,7 +103,8 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec)
|
||||
}
|
||||
}
|
||||
let dataset = table.dataset.get().await?;
|
||||
validate_maintained_indexes(&dataset, requested).await?;
|
||||
validate_maintained_indexes_with(&dataset, requested, &table.mem_index_registry())
|
||||
.await?;
|
||||
Some(requested.to_vec())
|
||||
}
|
||||
None => None,
|
||||
@@ -126,7 +127,9 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec)
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
let mut builder = dataset.initialize_mem_wal();
|
||||
let mut builder = dataset
|
||||
.initialize_mem_wal()
|
||||
.mem_index_registry(table.mem_index_registry());
|
||||
let writer_config_defaults = match spec {
|
||||
LsmWriteSpec::Bucket {
|
||||
column,
|
||||
@@ -725,7 +728,8 @@ pub(crate) async fn execute_lsm_merge_insert(
|
||||
return Ok(lsm_merge_result(0));
|
||||
};
|
||||
|
||||
let config = shard_writer_config_from_defaults(&plan.writer_config_defaults);
|
||||
let mut config = shard_writer_config_from_defaults(&plan.writer_config_defaults);
|
||||
config.mem_index_registry = table.mem_index_registry();
|
||||
let writer = table
|
||||
.dataset
|
||||
.shard_writer()
|
||||
|
||||
@@ -33,7 +33,6 @@ use datafusion_physical_plan::expressions::Column;
|
||||
use datafusion_physical_plan::projection::ProjectionExec;
|
||||
use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
|
||||
use lance::Dataset;
|
||||
use lance::dataset::mem_wal::index::{IndexStore, VectorMemQuery};
|
||||
use lance::dataset::mem_wal::scanner::{InMemoryMemTableRef, InMemoryMemTables};
|
||||
use lance::dataset::mem_wal::{
|
||||
DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig,
|
||||
@@ -711,17 +710,14 @@ async fn vector_plan(
|
||||
|
||||
// The base arm relies on the column's vector index (`fast_search`). Unmaintained,
|
||||
// its catch-up is untracked and exclusion falls back to the compaction watermark,
|
||||
// dropping compacted SSTables the lagging base index has not re-indexed; and a
|
||||
// resident MemTable built before the index joined the set carries none either.
|
||||
// dropping compacted SSTables the lagging base index has not re-indexed.
|
||||
// Reject rather than silently omit rows, mirroring the FTS arm.
|
||||
if !index_maintained(dataset, &column, details, "VectorIndexDetails").await?
|
||||
|| !resident_memtables_carry(&in_memory, |memtable| {
|
||||
carries_vector_index(&memtable.index_store, &column)
|
||||
})
|
||||
{
|
||||
// A resident MemTable needs no index of its own: one that cannot answer
|
||||
// compares every row.
|
||||
if !index_maintained(dataset, &column, details, "VectorIndexDetails").await? {
|
||||
return Err(Error::NotSupported {
|
||||
message: format!(
|
||||
"the MemWAL LSM scanner requires the vector index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes) and carried by every resident MemTable; otherwise compacted rows not yet re-indexed are omitted. set use_lsm(false) to read the base table only"
|
||||
"the MemWAL LSM scanner requires the vector index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes); otherwise compacted rows not yet re-indexed are omitted. set use_lsm(false) to read the base table only"
|
||||
),
|
||||
});
|
||||
}
|
||||
@@ -764,22 +760,6 @@ async fn vector_plan(
|
||||
Ok(scanner.create_plan().await?)
|
||||
}
|
||||
|
||||
/// Probe for a resident vector index without imposing a distance metric, matching
|
||||
/// the presence check used before Lance introduced query-based index lookup.
|
||||
fn carries_vector_index(store: &IndexStore, column: &str) -> bool {
|
||||
let probe = VectorMemQuery {
|
||||
vector: arrow_array::FixedSizeListArray::new_null(
|
||||
Arc::new(arrow_schema::Field::new("item", DataType::Float32, true)),
|
||||
1,
|
||||
1,
|
||||
),
|
||||
k: 1,
|
||||
ef: None,
|
||||
distance_type: None,
|
||||
};
|
||||
store.index_answering(column, &probe).is_some()
|
||||
}
|
||||
|
||||
/// Whether `column` stores binary (uint8) vectors, which the LSM vector planner
|
||||
/// does not support.
|
||||
fn is_binary_vector_column(schema: &ArrowSchema, column: &str) -> bool {
|
||||
@@ -838,9 +818,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resident_index_checks_distinguish_columns_and_query_types() {
|
||||
fn resident_fts_check_distinguishes_columns() {
|
||||
use arrow_schema::Field;
|
||||
use lance::dataset::mem_wal::index::MemIndexSpec;
|
||||
use lance::dataset::mem_wal::index::{IndexStore, MemIndexSpec};
|
||||
|
||||
let schema = ArrowSchema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
@@ -863,16 +843,11 @@ mod tests {
|
||||
),
|
||||
];
|
||||
let store = IndexStore::from_specs(&specs, &schema, 1_000, 16).unwrap();
|
||||
assert!(carries_vector_index(&store, "vector"));
|
||||
for column in ["id", "text", "missing"] {
|
||||
assert!(!carries_vector_index(&store, column));
|
||||
}
|
||||
assert!(!store.fts_granularities_on("text").is_empty());
|
||||
for column in ["id", "vector", "missing"] {
|
||||
assert!(store.fts_granularities_on(column).is_empty());
|
||||
}
|
||||
let empty = IndexStore::from_specs(&[], &schema, 1_000, 16).unwrap();
|
||||
assert!(!carries_vector_index(&empty, "vector"));
|
||||
assert!(empty.fts_granularities_on("text").is_empty());
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user