diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 784ee1fa76c..2a09106ed27 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -93,6 +93,7 @@ pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { } pub use api::{DatasetMemWalExt, InitializeMemWalBuilder}; +pub use index::{MemIndexKind, is_maintainable_index_type}; pub use manifest::ShardManifestStore; pub use memtable::scanner::MemTableScanner; pub use scanner::{LsmDataSource, LsmGeneration, LsmScanner, ShardSnapshot}; diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index 005b47d2b53..b927c32623a 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -24,6 +24,7 @@ use crate::index::DatasetIndexInternalExt; use crate::index::mem_wal::{load_mem_wal_index_details, new_mem_wal_index_meta}; use super::ShardWriterConfig; +use super::index::{MemIndexKind, unsupported_index_type}; use super::scanner::sstable_cache::open_sstable; use super::scanner::{DatasetCache, ShardSnapshot}; use super::util::derived_store_params; @@ -651,41 +652,39 @@ impl DatasetMemWalExt for Dataset { )) })?; - // Detect index type and create appropriate config + // Detect index kind and create appropriate config let type_url = index_meta .index_details .as_ref() .map(|d| d.type_url.as_str()) .unwrap_or(""); - let index_type = MemIndexConfig::detect_index_type(type_url)?; + let kind = MemIndexKind::from_type_url(type_url) + .ok_or_else(|| unsupported_index_type(type_url))?; - match index_type { - "btree" => { + // Exhaustive: a new kind must be built here, or callers filtering on + // `is_maintainable_index_type` would admit an index this writer + // cannot open, failing every memtable claim. + match kind { + MemIndexKind::BTree => { index_configs.push(MemIndexConfig::btree_from_metadata( &index_meta, self.schema(), )?); } - "fts" => { + MemIndexKind::Fts => { index_configs.push(MemIndexConfig::fts_from_metadata( &index_meta, self.schema(), )?); } - "vector" => { + MemIndexKind::Hnsw => { let hnsw_params = config.hnsw_params.get(index_name).cloned(); let vector_config = load_vector_index_config(self, index_name, &index_meta, hnsw_params) .await?; index_configs.push(vector_config); } - _ => { - return Err(Error::invalid_input(format!( - "Unknown index type: {}", - index_type - ))); - } }; } diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index b3da5b5b67e..7525b573d54 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -259,9 +259,53 @@ fn is_encodable_pk_type(data_type: &DataType) -> bool { ) } -/// Configuration for an index in MemWAL. +/// The index kinds a MemTable can maintain — the registry of MemWAL index +/// support. Data-free because indexes are identified by type url before any +/// [`MemIndexConfig`] exists. +/// +/// Adding a variant is a compile error in [`details_suffix`](Self::details_suffix), +/// `MemIndexConfig::kind`, and `Dataset::mem_wal_writer` until each handles it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MemIndexKind { + /// BTree index for scalar fields (point lookups, range queries). + BTree, + /// HNSW vector index built incrementally, queryable while building. + Hnsw, + /// Full-text search index. + Fts, +} + +impl MemIndexKind { + /// Every maintainable kind. A kind missing here is never detected, so it + /// goes unmaintained rather than reaching a memtable that cannot build it. + pub const ALL: &'static [Self] = &[Self::BTree, Self::Hnsw, Self::Fts]; + + /// Suffix of the protobuf details message identifying this kind. + /// + /// Only the suffix: the prefix varies by dataset version + /// (`/lance.table.`, `/lance.index.pb.`, and the `type.googleapis.com/` + /// form MemWAL flush once wrote), and all must resolve. + pub const fn details_suffix(self) -> &'static str { + match self { + Self::BTree => "BTreeIndexDetails", + Self::Hnsw => "VectorIndexDetails", + Self::Fts => "InvertedIndexDetails", + } + } + + /// The kind a base-table index of this protobuf type maps to, or `None` + /// when a memtable cannot maintain it. + pub fn from_type_url(type_url: &str) -> Option { + Self::ALL + .iter() + .copied() + .find(|kind| type_url.ends_with(kind.details_suffix())) + } +} + +/// Configuration for an index in MemWAL. Pairs 1:1 with [`MemIndexKind`] via +/// [`kind`](Self::kind). /// -/// Each variant contains all the configuration needed for that index type. /// `Hnsw` is boxed because `HnswBuildParams` is small but the variant may /// grow with future config (e.g. shard-specific tuning). #[derive(Debug, Clone)] @@ -275,6 +319,16 @@ pub enum MemIndexConfig { } impl MemIndexConfig { + /// The kind this config builds. Links the config enum to the registry, so + /// a new variant must declare its kind. + pub const fn kind(&self) -> MemIndexKind { + match self { + Self::BTree(_) => MemIndexKind::BTree, + Self::Hnsw(_) => MemIndexKind::Hnsw, + Self::Fts(_) => MemIndexKind::Fts, + } + } + /// Get the index name. pub fn name(&self) -> &str { match self { @@ -363,22 +417,6 @@ impl MemIndexConfig { )) } - /// Detect index type from protobuf type_url. - pub fn detect_index_type(type_url: &str) -> Result<&'static str> { - if type_url.ends_with("BTreeIndexDetails") { - Ok("btree") - } else if type_url.ends_with("InvertedIndexDetails") { - Ok("fts") - } else if type_url.ends_with("VectorIndexDetails") { - Ok("vector") - } else { - Err(Error::invalid_input(format!( - "Unsupported index type for MemWAL: {}. Supported: BTree, Inverted, Vector", - type_url - ))) - } - } - fn fts_format_version_from_metadata( index_meta: &IndexMetadata, ) -> Result { @@ -415,6 +453,23 @@ impl MemIndexConfig { } } +/// Whether the MemWAL can maintain an index of this protobuf type. +/// +/// Opening a shard writer rejects anything outside this set, which makes the +/// table unwritable — so filter on this before committing a maintained set, +/// not at claim time. +pub fn is_maintainable_index_type(type_url: &str) -> bool { + MemIndexKind::from_type_url(type_url).is_some() +} + +/// Shared by the detection and writer paths so both report the same thing. +pub(crate) fn unsupported_index_type(type_url: &str) -> Error { + Error::invalid_input(format!( + "Unsupported index type for MemWAL: {}. Supported: BTree, Inverted, Vector", + type_url + )) +} + /// Registry managing all in-memory indexes for a MemTable. /// /// Indexes are keyed by index name. Each index stores its field_id for @@ -1172,26 +1227,59 @@ mod tests { use super::*; use arrow_array::{Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; - use log::warn; use rstest::rstest; use std::sync::Arc; use uuid::Uuid; - /// Check if an index type is supported and log warning if not. - fn check_index_type_supported(index_type: &str) -> bool { - match index_type.to_lowercase().as_str() { - "btree" | "scalar" => true, - "hnsw" | "vector" => true, - "fts" | "inverted" | "fulltext" => true, - _ => { - warn!( - "Index type '{}' is not supported for MemWAL. \ - Supported types: btree, hnsw, fts. Skipping.", - index_type - ); - false - } + /// Matching is on the message-name suffix, not the whole url: `Any::from_msg` + /// emits the package (`/lance.table.`, `/lance.index.pb.`), while MemWAL flush + /// used to hand-write a `type.googleapis.com/` url that existing datasets + /// still carry. + #[rstest] + #[case::btree("/lance.table.BTreeIndexDetails", Some(MemIndexKind::BTree))] + #[case::fts("/lance.table.InvertedIndexDetails", Some(MemIndexKind::Fts))] + #[case::fts_legacy("/lance.index.pb.InvertedIndexDetails", Some(MemIndexKind::Fts))] + #[case::vector("/lance.index.pb.VectorIndexDetails", Some(MemIndexKind::Hnsw))] + // What MemWAL flush wrote before it switched to `Any::from_msg`. + #[case::vector_legacy_flush( + "type.googleapis.com/lance.index.VectorIndexDetails", + Some(MemIndexKind::Hnsw) + )] + #[case::bitmap("/lance.table.BitmapIndexDetails", None)] + #[case::label_list("/lance.table.LabelListIndexDetails", None)] + #[case::ngram("/lance.table.NGramIndexDetails", None)] + #[case::zone_map("/lance.table.ZoneMapIndexDetails", None)] + #[case::bloom_filter("/lance.index.pb.BloomFilterIndexDetails", None)] + #[case::json("/lance.index.pb.JsonIndexDetails", None)] + #[case::fm("/lance.index.pb.FMIndexDetails", None)] + #[case::absent("", None)] + fn type_urls_resolve_to_the_kind_the_writer_builds( + #[case] type_url: &str, + #[case] expected: Option, + ) { + assert_eq!(MemIndexKind::from_type_url(type_url), expected); + assert_eq!(is_maintainable_index_type(type_url), expected.is_some()); + } + + /// `ALL` is hand-maintained, so a kind left out of it stops resolving. + #[test] + fn every_kind_is_registered_and_uniquely_identified() { + for kind in MemIndexKind::ALL { + assert_eq!( + MemIndexKind::from_type_url(&format!("/lance.table.{}", kind.details_suffix())), + Some(*kind), + "{kind:?} does not resolve from its own suffix", + ); } + let suffixes: std::collections::HashSet<_> = MemIndexKind::ALL + .iter() + .map(|k| k.details_suffix()) + .collect(); + assert_eq!( + suffixes.len(), + MemIndexKind::ALL.len(), + "two kinds share a details suffix, so one can never be resolved", + ); } fn create_test_schema() -> Arc { @@ -1546,18 +1634,6 @@ mod tests { assert_eq!(fts.doc_count(), 3); } - #[test] - fn test_check_index_type_supported() { - assert!(check_index_type_supported("btree")); - assert!(check_index_type_supported("BTree")); - assert!(check_index_type_supported("hnsw")); - assert!(check_index_type_supported("vector")); - assert!(check_index_type_supported("fts")); - assert!(check_index_type_supported("inverted")); - - assert!(!check_index_type_supported("unknown")); - } - #[test] fn fts_from_metadata_preserves_format_version() { let arrow_schema = create_test_schema(); diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs index 0a5c97ee6c7..694d7f282ab 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/flush.rs @@ -33,6 +33,7 @@ use crate::dataset::mem_wal::manifest::ShardManifestStore; use crate::dataset::mem_wal::scanner::SsTableWarmer; use crate::dataset::mem_wal::scanner::exec::{compute_pk_hash, validate_pk_types}; use crate::dataset::mem_wal::util::{derived_store_params, generate_random_hash, sstable_path}; +use crate::index::vector::details::vector_index_details_default; use crate::session::Session; #[derive(Debug, Clone)] @@ -1105,10 +1106,9 @@ impl MemTableFlusher { ); index_writer.finish().await?; - let index_details = Some(Arc::new(prost_types::Any { - type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), - value: vec![], - })); + // Packed the same way index creation does; hand-building the `Any` here + // produced a `type.googleapis.com/` url no other writer in lance emits. + let index_details = Some(Arc::new(vector_index_details_default())); let index_meta = IndexMetadata { uuid: index_uuid, name: config.name.clone(), diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 7f05cbfff3e..3e588ca35ce 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -37,7 +37,7 @@ use uuid::Uuid; pub use super::index::{ BTreeIndexConfig, BTreeMemIndex, FtsIndexConfig, HnswIndexConfig, IndexStore, MemIndexConfig, - validate_index_configs, + MemIndexKind, validate_index_configs, }; pub use super::memtable::CacheConfig; pub use super::memtable::MemTable;