From 02490a4156c86cad70ce53cc40404b53f9c095c5 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 15 Jul 2026 02:41:18 +0800 Subject: [PATCH 1/8] feat(index): support element-document FTS targets --- protos/index_old.proto | 12 + python/python/lance/dataset.py | 11 +- python/python/tests/test_scalar_index.py | 83 ++ rust/lance-core/src/datatypes.rs | 4 +- rust/lance-core/src/datatypes/schema.rs | 175 +++- rust/lance-index/src/scalar/inverted.rs | 31 +- .../src/scalar/inverted/builder.rs | 75 +- rust/lance-index/src/scalar/inverted/index.rs | 610 ++++++++++++-- .../src/scalar/inverted/lazy_docset.rs | 77 +- .../src/scalar/inverted/tokenizer.rs | 94 ++- rust/lance-index/src/scalar/inverted/wand.rs | 3 + rust/lance/src/dataset/mem_wal/index.rs | 104 ++- rust/lance/src/dataset/mem_wal/index/fts.rs | 692 ++++++++++++---- .../src/dataset/mem_wal/memtable/flush.rs | 11 +- .../mem_wal/memtable/scanner/builder.rs | 11 +- .../mem_wal/memtable/scanner/exec/fts.rs | 212 ++++- .../src/dataset/mem_wal/scanner/fts_search.rs | 260 +++++- rust/lance/src/dataset/scanner.rs | 331 ++++++-- rust/lance/src/index/create.rs | 93 ++- rust/lance/src/index/scalar.rs | 83 +- rust/lance/src/index/scalar/inverted.rs | 255 +++++- rust/lance/src/io/exec/fts.rs | 770 ++++++++++++++---- rust/lance/tests/query/inverted.rs | 669 ++++++++++++++- 23 files changed, 4061 insertions(+), 605 deletions(-) diff --git a/protos/index_old.proto b/protos/index_old.proto index eb984d6fe29..695a9265a24 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -25,6 +25,12 @@ message BitmapIndexDetails {} message LabelListIndexDetails {} message NGramIndexDetails {} message ZoneMapIndexDetails {} +message FtsTarget { + // The field id of the text value that is tokenized. + int32 source_field_id = 1; + // Ordered list boundaries crossed between the public field and source value. + repeated int32 boundary_field_ids = 2; +} message InvertedIndexDetails { // Marking this field as optional as old versions of the index store blank details and we // need to make sure we have a proper optional field to detect this. @@ -44,4 +50,10 @@ message InvertedIndexDetails { // A present value records the block size used by the index; 256 is only // valid with format version 3. optional uint32 block_size = 12; + // The logical FTS document target. An absent value denotes a legacy + // row-document index and is normalized from IndexMetadata.fields. + FtsTarget fts_target = 13; + // The posting-list payload format. This is separate from index_version, + // which is an element-document compatibility marker for version 4. + optional uint32 posting_format_version = 14; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index a1ce77b82b8..fe6cb212cea 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3124,8 +3124,6 @@ def _prepare_scalar_index_request( column = column[0] lance_field = self._ds.lance_schema.field_case_insensitive(column) - if lance_field is None: - raise KeyError(f"{column} not found in schema") if isinstance(index_type, str): index_type = index_type.upper() @@ -3148,6 +3146,13 @@ def _prepare_scalar_index_request( ) ) + if lance_field is None: + if index_type in ["INVERTED", "FTS"]: + # Public FTS target paths such as `tags[*]` are resolved + # against stable field ids by the Rust implementation. + return column, index_type, index_type + raise KeyError(f"{column} not found in schema") + field = lance_field.to_arrow() field_type = field.type @@ -3201,6 +3206,8 @@ def _prepare_scalar_index_request( return column, index_type, index_type elif isinstance(index_type, IndexConfig): logical_index_type = index_type.index_type.upper() + if lance_field is None and logical_index_type not in ["INVERTED", "FTS"]: + raise KeyError(f"{column} not found in schema") config = json.dumps(index_type.parameters) kwargs["config"] = indices.IndexConfig(index_type.index_type, config) return column, "scalar", logical_index_type diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 985072fe5de..88e00e0e0d5 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -1351,6 +1351,89 @@ def test_fts_on_list(tmp_path): assert results.num_rows == 6 +@pytest.mark.parametrize( + "list_type", + [ + pa.list_(pa.string()), + pa.list_(pa.large_string()), + pa.large_list(pa.string()), + pa.large_list(pa.large_string()), + ], +) +def test_fts_on_list_elements(tmp_path, list_type): + data = pa.table( + { + "id": pa.array([0, 1, 2]), + "tags": pa.array( + [ + ["alpha beta", "gamma alpha", None, "", "delta"], + ["beta", "gamma"], + None, + ], + type=list_type, + ), + } + ) + ds = lance.write_dataset(data, tmp_path) + + def hits(table): + return sorted(zip(table["id"].to_pylist(), table["_doc_index"].to_pylist())) + + query = MatchQuery("alpha", "tags[*]") + flat = ds.to_table(full_text_query=query) + assert hits(flat) == [(0, [0]), (0, [1])] + assert pa.types.is_list(flat.schema.field("_doc_index").type) + assert flat.schema.field("_doc_index").type.value_type == pa.uint32() + assert hits(ds.to_table(full_text_query=MatchQuery("delta", "tags[*]"))) == [ + (0, [4]) + ] + + ds.create_scalar_index("tags", "INVERTED", with_position=True) + ds.create_scalar_index("tags[*]", "INVERTED", with_position=True) + ds.create_scalar_index( + "tags[*]", + IndexConfig(index_type="inverted", parameters={"with_position": True}), + ) + index_names = {index.name for index in ds.describe_indices()} + assert {"tags_idx", "tags[*]_idx"}.issubset(index_names) + row_auto = ds.to_table(full_text_query="alpha") + assert row_auto["id"].to_pylist() == [0] + assert "_doc_index" not in row_auto.column_names + indexed = ds.to_table(full_text_query=query) + assert hits(indexed) == [(0, [0]), (0, [1])] + assert pa.types.is_list(indexed.schema.field("_doc_index").type) + assert indexed.schema.field("_doc_index").type.value_type == pa.uint32() + filtered = ds.to_table(full_text_query=query, filter="id = 0", prefilter=True) + assert hits(filtered) == [(0, [0]), (0, [1])] + assert hits(ds.to_table(full_text_query=MatchQuery("delta", "tags[*]"))) == [ + (0, [4]) + ] + + phrase = ds.to_table(full_text_query=PhraseQuery("beta gamma", "tags[*]")) + assert phrase.num_rows == 0 + assert hits(ds.to_table(full_text_query=PhraseQuery("alpha beta", "tags[*]"))) == [ + (0, [0]) + ] + row_phrase = ds.to_table(full_text_query=PhraseQuery("beta gamma", "tags")) + assert sorted(row_phrase["id"].to_pylist()) == [0, 1] + assert "_doc_index" not in row_phrase.column_names + + ds.insert( + pa.table( + { + "id": pa.array([3]), + "tags": pa.array([["alpha", "alpha again"]], type=list_type), + } + ) + ) + assert hits(ds.to_table(full_text_query=query)) == [ + (0, [0]), + (0, [1]), + (3, [0]), + (3, [1]), + ] + + def test_fts_fuzzy_query(tmp_path): data = pa.table( { diff --git a/rust/lance-core/src/datatypes.rs b/rust/lance-core/src/datatypes.rs index 2f5c7b0680e..86ec338cf3b 100644 --- a/rust/lance-core/src/datatypes.rs +++ b/rust/lance-core/src/datatypes.rs @@ -23,9 +23,9 @@ pub use field::{ OnTypeMismatch, SchemaCompareOptions, }; pub use schema::{ - BlobHandling, FieldRef, OnMissing, Projectable, Projection, Schema, + BlobHandling, FieldPathComponent, FieldRef, OnMissing, Projectable, Projection, Schema, escape_field_path_for_project, format_field_path, parse_field_path, - validate_fixed_size_list_dimensions, + parse_field_path_components, validate_fixed_size_list_dimensions, }; pub static BLOB_DESC_FIELDS: LazyLock = LazyLock::new(|| { diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 6f9fc61b334..da3f2d0b7e1 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -1505,21 +1505,21 @@ impl Projection { } } -/// Parse a field path that may contain quoted field names. -/// -/// Field names containing dots must be quoted with backticks. -/// For example: "parent.`child.with.dot`" parses to ["parent", "child.with.dot"] -/// -/// Backticks within quoted fields must be escaped by doubling them. -/// For example: "`field``with``backticks`" represents the field name "field`with`backticks" -/// -/// Returns an error if: -/// - The input path is empty -/// - The path has malformed quotes (unclosed, misplaced, etc.) -/// - The path has empty segments (e.g., "parent..child" or "parent.") +/// A component of a parsed field path. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FieldPathComponent { + /// A named field. + Field(String), + /// The elements of a list field, written as `[*]`. + ListWildcard, +} + +/// Parse a field path that may contain quoted field names and list wildcards. /// -/// The result is guaranteed to contain at least one element. -pub fn parse_field_path(path: &str) -> Result> { +/// Unquoted `[*]` is parsed as [`FieldPathComponent::ListWildcard`]. To address +/// a field whose name literally contains `[*]`, quote the field name with +/// backticks. +pub fn parse_field_path_components(path: &str) -> Result> { if path.is_empty() { return Err(Error::schema("Field path cannot be empty".to_string())); } @@ -1527,35 +1527,25 @@ pub fn parse_field_path(path: &str) -> Result> { let mut result = Vec::new(); let mut current = String::new(); let mut in_quotes = false; + let mut just_closed_quote = false; + let mut trailing_dot = false; let mut chars = path.chars().peekable(); while let Some(ch) = chars.next() { match ch { '`' => { if in_quotes { - // Check if this is an escaped backtick (double backtick) if chars.peek() == Some(&'`') { - // Consume the second backtick and add a single backtick to current chars.next(); current.push('`'); } else { - // End of quoted field in_quotes = false; - // After closing quote, we should either see a dot or end of string - if let Some(&next_ch) = chars.peek() - && next_ch != '.' - { - return Err(Error::schema(format!( - "Invalid field path '{}': expected '.' or end of string after closing quote", - path - ))); - } + just_closed_quote = true; } - } else if current.is_empty() { - // Start of quoted field + } else if current.is_empty() && !just_closed_quote { in_quotes = true; + trailing_dot = false; } else { - // Quote in the middle of unquoted field name return Err(Error::schema(format!( "Invalid field path '{}': unexpected quote in the middle of field name", path @@ -1563,17 +1553,75 @@ pub fn parse_field_path(path: &str) -> Result> { } } '.' if !in_quotes => { - if current.is_empty() { + if trailing_dot { return Err(Error::schema(format!( "Invalid field path '{}': empty field name", path ))); } - result.push(current); - current = String::new(); + if !current.is_empty() { + result.push(FieldPathComponent::Field(std::mem::take(&mut current))); + } else if !matches!(result.last(), Some(FieldPathComponent::ListWildcard)) { + return Err(Error::schema(format!( + "Invalid field path '{}': empty field name", + path + ))); + } + just_closed_quote = false; + trailing_dot = true; + } + '[' if !in_quotes && chars.peek() == Some(&'*') => { + if trailing_dot { + return Err(Error::schema(format!( + "Invalid field path '{}': list wildcard must immediately follow a field or list wildcard", + path + ))); + } + let mut lookahead = chars.clone(); + lookahead.next(); + if lookahead.next() == Some(']') { + if !current.is_empty() { + result.push(FieldPathComponent::Field(std::mem::take(&mut current))); + } else if !matches!(result.last(), Some(FieldPathComponent::ListWildcard)) { + return Err(Error::schema(format!( + "Invalid field path '{}': list wildcard must follow a field", + path + ))); + } + chars.next(); + chars.next(); + result.push(FieldPathComponent::ListWildcard); + just_closed_quote = false; + trailing_dot = false; + if let Some(next) = chars.peek() + && *next != '.' + && *next != '[' + { + return Err(Error::schema(format!( + "Invalid field path '{}': expected '.', '[*]', or end of string after list wildcard", + path + ))); + } + } else { + if just_closed_quote { + return Err(Error::schema(format!( + "Invalid field path '{}': expected '.', '[*]', or end of string after closing quote", + path + ))); + } + current.push(ch); + trailing_dot = false; + } } _ => { + if !in_quotes && just_closed_quote { + return Err(Error::schema(format!( + "Invalid field path '{}': expected '.', '[*]', or end of string after closing quote", + path + ))); + } current.push(ch); + trailing_dot = false; } } } @@ -1584,25 +1632,47 @@ pub fn parse_field_path(path: &str) -> Result> { path ))); } - - if !current.is_empty() { - result.push(current); - } else if !result.is_empty() { + if trailing_dot { return Err(Error::schema(format!( "Invalid field path '{}': trailing dot", path ))); } - - // This check is now redundant since we check for empty input at the beginning, - // but keeping it for extra safety + if !current.is_empty() { + result.push(FieldPathComponent::Field(current)); + } if result.is_empty() { return Err(Error::schema(format!("Invalid field path '{}'", path))); } - Ok(result) } +/// Parse a field path that may contain quoted field names. +/// +/// Field names containing dots must be quoted with backticks. +/// For example: "parent.`child.with.dot`" parses to ["parent", "child.with.dot"] +/// +/// Backticks within quoted fields must be escaped by doubling them. +/// For example: "`field``with``backticks`" represents the field name "field`with`backticks" +/// +/// Returns an error if: +/// - The input path is empty +/// - The path has malformed quotes (unclosed, misplaced, etc.) +/// - The path has empty segments (e.g., "parent..child" or "parent.") +/// +/// The result is guaranteed to contain at least one element. +pub fn parse_field_path(path: &str) -> Result> { + parse_field_path_components(path).map(|components| { + components + .into_iter() + .map(|component| match component { + FieldPathComponent::Field(field) => field, + FieldPathComponent::ListWildcard => "[*]".to_string(), + }) + .collect() + }) +} + /// Format a field path, quoting field names that require escaping. /// /// Field names are quoted if they contain any character that is not alphanumeric @@ -1795,6 +1865,31 @@ mod tests { // Invalid: trailing dot assert!(parse_field_path("parent.").is_err()); + assert_eq!( + parse_field_path_components("tags[*]").unwrap(), + vec![ + FieldPathComponent::Field("tags".to_string()), + FieldPathComponent::ListWildcard, + ] + ); + assert_eq!( + parse_field_path_components("parent.items[*].text").unwrap(), + vec![ + FieldPathComponent::Field("parent".to_string()), + FieldPathComponent::Field("items".to_string()), + FieldPathComponent::ListWildcard, + FieldPathComponent::Field("text".to_string()), + ] + ); + assert_eq!( + parse_field_path_components("`tags[*]`").unwrap(), + vec![FieldPathComponent::Field("tags[*]".to_string())] + ); + assert!(parse_field_path_components("[*]").is_err()); + assert!(parse_field_path_components("tags[*]suffix").is_err()); + assert!(parse_field_path_components("tags[*]..value").is_err()); + assert!(parse_field_path_components("tags[*].[*]").is_err()); + // Test formatting assert_eq!( format_field_path(&["parent", "child.with.dot", "normal"]), diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 926bfad6a69..f1720931246 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -149,6 +149,9 @@ impl InvertedIndexPlugin { params.validate_format_version()?; let format_version = params.resolved_format_version(); + let is_element_document = params + .fts_target() + .is_some_and(FtsTarget::is_element_document); let details = pbold::InvertedIndexDetails::try_from(¶ms)?; let mut inverted_index = InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask) @@ -156,7 +159,11 @@ impl InvertedIndexPlugin { let files = inverted_index.update(data, index_store, None).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: format_version.index_version(), + index_version: if is_element_document { + INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT + } else { + format_version.index_version() + }, files, }) } @@ -268,7 +275,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { } fn version(&self) -> u32 { - max_supported_fts_format_version().index_version() + INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT } fn new_query_parser( @@ -301,10 +308,19 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { - Ok( + let details = _index_details.to_msg::()?; + let index = if let Some(target) = details.fts_target.as_ref() { + InvertedIndex::load_with_target( + index_store, + frag_reuse_index, + cache, + FtsTarget::from(target), + ) + .await? + } else { InvertedIndex::load(index_store, frag_reuse_index, cache).await? - as Arc, - ) + }; + Ok(index as Arc) } fn details_as_json(&self, details: &prost_types::Any) -> Result { @@ -322,10 +338,7 @@ mod tests { #[test] fn test_plugin_version_tracks_max_supported_format() { let plugin = InvertedIndexPlugin; - assert_eq!( - plugin.version(), - max_supported_fts_format_version().index_version() - ); + assert_eq!(plugin.version(), INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT); } #[test] diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 41804e94fd1..7b05a47427f 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -388,6 +388,11 @@ impl InvertedIndexBuilder { token_set_format: self.token_set_format, worker_memory_limit_bytes, block_size: self.params.block_size, + coordinate_rank: self + .params + .fts_target() + .map(|target| target.coordinate_rank()) + .unwrap_or(0), }; let next_id = self.next_partition_id(); let id_alloc = Arc::new(AtomicU64::new(next_id)); @@ -512,6 +517,10 @@ impl InvertedIndexBuilder { None, &LanceCache::no_cache(), self.token_set_format, + self.params + .fts_target() + .map(|target| target.coordinate_rank()) + .unwrap_or(0), ) .await?; let mut builder = part.into_builder().await?; @@ -1005,8 +1014,16 @@ impl InnerBuilder { } let doc_id_offset = self.docs.len() as u32; - for (row_id, num_tokens) in docs.iter() { - self.docs.append(*row_id, *num_tokens); + for doc_id in 0..docs.len() as u32 { + let row_id = docs.row_id(doc_id); + let num_tokens = docs.num_tokens(doc_id); + let doc_index = docs.doc_index(doc_id); + if doc_index.is_empty() { + self.docs.append(row_id, num_tokens); + } else { + self.docs + .append_with_doc_index(row_id, num_tokens, &doc_index)?; + } } self.posting_lists.resize_with(self.tokens.len(), || { PostingListBuilder::new_with_posting_tail_codec_and_block_size( @@ -1246,6 +1263,7 @@ struct IndexWorker { token_set_format: TokenSetFormat, token_ids: Vec, last_token_count: usize, + coordinate_rank: usize, } struct TailPartition { @@ -1271,6 +1289,7 @@ struct IndexWorkerConfig { token_set_format: TokenSetFormat, worker_memory_limit_bytes: u64, block_size: usize, + coordinate_rank: usize, } impl IndexWorker { @@ -1342,6 +1361,7 @@ impl IndexWorker { token_set_format: config.token_set_format, token_ids: Vec::new(), last_token_count: 0, + coordinate_rank: config.coordinate_rank, }) } @@ -1362,7 +1382,7 @@ impl IndexWorker { .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); for (doc, row_id) in docs { - self.process_document(row_id, DocumentSource::Text(doc), false) + self.process_document(row_id, DocumentSource::Text(doc), &[], false) .await?; } } @@ -1406,8 +1426,29 @@ impl IndexWorker { continue; }; - self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), true) - .await?; + if self.coordinate_rank == 0 { + self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), &[], true) + .await?; + } else { + if self.coordinate_rank != 1 { + return Err(Error::not_supported(format!( + "FTS element documents currently support exactly one list boundary, got {}", + self.coordinate_rank + ))); + } + for (ordinal, element) in iter_str_array(doc.as_ref()).enumerate() { + let ordinal = u32::try_from(ordinal).map_err(|_| { + Error::invalid_input(format!( + "FTS element ordinal overflow for row_id={row_id}" + )) + })?; + let Some(element) = element else { + continue; + }; + self.process_document(*row_id, DocumentSource::Text(element), &[ordinal], true) + .await?; + } + } } Ok(()) @@ -1436,6 +1477,7 @@ impl IndexWorker { &mut self, row_id: u64, document: DocumentSource<'_>, + doc_index: &[u32], skip_empty_document: bool, ) -> Result<()> { let with_position = self.has_position(); @@ -1575,7 +1617,13 @@ impl IndexWorker { } let old_doc_memory_size = self.builder.docs.memory_size() as u64; - let appended_doc_id = self.builder.docs.append(row_id, token_num); + let appended_doc_id = if doc_index.is_empty() { + self.builder.docs.append(row_id, token_num) + } else { + self.builder + .docs + .append_with_doc_index(row_id, token_num, doc_index)? + }; debug_assert_eq!(appended_doc_id, doc_id); self.adjust_tracked_memory_size( old_doc_memory_size, @@ -1742,6 +1790,7 @@ impl PositionRecorder { #[derive(Debug, Eq, PartialEq, Clone, DeepSizeOf)] pub struct ScoredDoc { pub row_id: u64, + pub doc_index: Vec, pub score: OrderedFloat, } @@ -1749,6 +1798,15 @@ impl ScoredDoc { pub fn new(row_id: u64, score: f32) -> Self { Self { row_id, + doc_index: Vec::new(), + score: OrderedFloat(score), + } + } + + pub fn with_doc_index(row_id: u64, doc_index: Vec, score: f32) -> Self { + Self { + row_id, + doc_index, score: OrderedFloat(score), } } @@ -3274,6 +3332,7 @@ mod tests { token_set_format, worker_memory_limit_bytes: u64::MAX, block_size: params.block_size, + coordinate_rank: 0, }, ) .await?; @@ -3298,6 +3357,7 @@ mod tests { token_set_format, worker_memory_limit_bytes: u64::MAX, block_size: params.block_size, + coordinate_rank: 0, }, ) .await?; @@ -3781,6 +3841,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; @@ -3813,6 +3874,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; @@ -3852,6 +3914,7 @@ mod tests { token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, block_size: InvertedIndexParams::default().block_size, + coordinate_rank: 0, }, ) .await?; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 21aadf7d475..15c2f1cd87e 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -58,7 +58,7 @@ use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder} use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; -use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; +use super::{FtsTarget, InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ BLOCK_SIZE, ScoredDoc, doc_file_path, @@ -93,6 +93,8 @@ use std::str::FromStr; pub const INVERTED_INDEX_VERSION_V1: u32 = 1; pub const INVERTED_INDEX_VERSION_V2: u32 = 2; pub const INVERTED_INDEX_VERSION_V3: u32 = 3; +/// Element-document indexes add persisted document coordinates. +pub const INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT: u32 = 4; pub const TOKENS_FILE: &str = "tokens.lance"; pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; @@ -113,6 +115,8 @@ pub const MAX_SCORE_COL: &str = "_max_score"; pub const LENGTH_COL: &str = "_length"; pub const BLOCK_MAX_SCORE_COL: &str = "_block_max_score"; pub const NUM_TOKEN_COL: &str = "_num_tokens"; +pub const DOC_INDEX_COL: &str = "_doc_index"; +pub const DOC_INDEX_STORAGE_PREFIX: &str = "_doc_index_"; pub const SCORE_COL: &str = "_score"; pub const TOKEN_SET_FORMAT_KEY: &str = "token_set_format"; pub const POSTING_TAIL_CODEC_KEY: &str = "posting_tail_codec"; @@ -132,8 +136,30 @@ pub const ESTIMATED_MAX_TOKENS_PER_ROW: usize = 4 * 1024; pub static SCORE_FIELD: LazyLock = LazyLock::new(|| Field::new(SCORE_COL, DataType::Float32, true)); +pub static DOC_INDEX_FIELD: LazyLock = LazyLock::new(|| { + Field::new( + DOC_INDEX_COL, + DataType::List(Arc::new(Field::new("item", DataType::UInt32, false))), + false, + ) +}); pub static FTS_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(Schema::new(vec![ROW_ID_FIELD.clone(), SCORE_FIELD.clone()]))); +pub static ELEMENT_FTS_SCHEMA: LazyLock = LazyLock::new(|| { + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + DOC_INDEX_FIELD.clone(), + SCORE_FIELD.clone(), + ])) +}); + +pub fn fts_schema(target: &FtsTarget) -> SchemaRef { + if target.is_element_document() { + ELEMENT_FTS_SCHEMA.clone() + } else { + FTS_SCHEMA.clone() + } +} static ROW_ID_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(Schema::new(vec![ROW_ID_FIELD.clone()]))); @@ -551,7 +577,33 @@ impl DeepSizeOf for InvertedIndex { async fn resolve_deferred_candidates( docs: &LazyDocSet, candidates: &mut [DocCandidate], + coordinate_rank: usize, ) -> Result<()> { + if coordinate_rank > 0 { + let doc_ids = candidates + .iter() + .map(|candidate| { + u32::try_from(candidate.posting_doc_id).map_err(|_| { + Error::internal(format!( + "FTS partition-local doc id {} exceeds u32", + candidate.posting_doc_id + )) + }) + }) + .collect::>>()?; + let keys = docs.resolve_document_keys(&doc_ids).await?; + if keys.len() != candidates.len() { + return Err(Error::internal( + "resolve_document_keys returned an unexpected number of items".to_string(), + )); + } + for (candidate, (row_id, doc_index)) in candidates.iter_mut().zip(keys) { + candidate.addr = CandidateAddr::RowId(row_id); + candidate.doc_index = doc_index; + } + return Ok(()); + } + let pending: Vec = candidates .iter() .filter_map(|c| match c.addr { @@ -580,6 +632,13 @@ impl InvertedIndex { } fn index_version(&self) -> u32 { + if self + .params + .fts_target() + .is_some_and(FtsTarget::is_element_document) + { + return INVERTED_INDEX_VERSION_ELEMENT_DOCUMENT; + } match (self.token_set_format, self.format_version()) { ( TokenSetFormat::Arrow, @@ -882,6 +941,26 @@ impl InvertedIndex { metrics: Arc, base_scorer: Option<&MemBM25Scorer>, ) -> Result<(Vec, Vec)> { + let documents = self + .bm25_search_documents(tokens, params, operator, prefilter, metrics, base_scorer) + .await?; + Ok(documents + .into_iter() + .map(|document| (document.row_id, document.score.0)) + .unzip()) + } + + /// Search logical FTS documents, retaining element coordinates when present. + #[instrument(level = "debug", skip_all)] + pub async fn bm25_search_documents( + &self, + tokens: Arc, + params: Arc, + operator: Operator, + prefilter: Arc, + metrics: Arc, + base_scorer: Option<&MemBM25Scorer>, + ) -> Result> { // Fuzzy expansion runs once here, with the global `max_expansions` // budget, instead of once per partition: partitions receive the // final token list, so the matched terms cannot depend on how the @@ -896,7 +975,7 @@ impl InvertedIndex { .map(|idx| expanded.position(idx)) .collect::>(); if (0..tokens.len()).any(|idx| !surviving.contains(&tokens.position(idx))) { - return Ok((Vec::new(), Vec::new())); + return Ok(Vec::new()); } } expanded @@ -919,13 +998,14 @@ impl InvertedIndex { let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { - return Ok((Vec::new(), Vec::new())); + return Ok(Vec::new()); } fn push_scored_candidate( candidates: &mut BinaryHeap>, limit: usize, addr: CandidateAddr, + doc_index: Vec, score: f32, ) -> Result<()> { // resolve_deferred_candidates ran upstream, so every candidate @@ -940,10 +1020,10 @@ impl InvertedIndex { }; if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(ScoredDoc::with_doc_index(row_id, doc_index, score))); } else if candidates.peek().unwrap().0.score.0 < score { candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(ScoredDoc::with_doc_index(row_id, doc_index, score))); } Ok(()) } @@ -957,6 +1037,11 @@ impl InvertedIndex { // global k-th - see `Wand::shared_threshold`). let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let coordinate_rank = self + .params + .fts_target() + .map(FtsTarget::coordinate_rank) + .unwrap_or(0); let parts = self .partitions .iter() @@ -1042,8 +1127,12 @@ impl InvertedIndex { grouped_expansions, candidates, }; - resolve_deferred_candidates(&part.docs, &mut partition_result.candidates) - .await?; + resolve_deferred_candidates( + &part.docs, + &mut partition_result.candidates, + coordinate_rank, + ) + .await?; Result::Ok(partition_result) } }) @@ -1075,6 +1164,7 @@ impl InvertedIndex { if grouped_expansions.is_empty() { for DocCandidate { addr, + doc_index, freqs, doc_length, .. @@ -1086,7 +1176,7 @@ impl InvertedIndex { score += idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, addr, doc_index, score)?; } } else { let grouped_positions = grouped_expansions @@ -1095,6 +1185,7 @@ impl InvertedIndex { .collect::>(); for DocCandidate { addr, + doc_index, posting_doc_id, freqs, doc_length, @@ -1125,7 +1216,7 @@ impl InvertedIndex { score += idf_weight * scorer.doc_weight(freq, doc_length); } } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, addr, doc_index, score)?; } } } @@ -1133,8 +1224,8 @@ impl InvertedIndex { Ok(candidates .into_sorted_vec() .into_iter() - .map(|Reverse(doc)| (doc.row_id, doc.score.0)) - .unzip()) + .map(|Reverse(document)| document) + .collect()) } async fn load_legacy_index( @@ -1211,6 +1302,30 @@ impl InvertedIndex { frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> + where + Self: Sized, + { + Self::load_with_target_override(store, frag_reuse_index, index_cache, None).await + } + + pub async fn load_with_target( + store: Arc, + frag_reuse_index: Option>, + index_cache: &LanceCache, + target: FtsTarget, + ) -> Result> + where + Self: Sized, + { + Self::load_with_target_override(store, frag_reuse_index, index_cache, Some(target)).await + } + + async fn load_with_target_override( + store: Arc, + frag_reuse_index: Option>, + index_cache: &LanceCache, + target: Option, + ) -> Result> where Self: Sized, { @@ -1225,7 +1340,14 @@ impl InvertedIndex { .metadata .get("params") .ok_or(Error::index("params not found in metadata".to_owned()))?; - let params = serde_json::from_str::(params)?; + let mut params = serde_json::from_str::(params)?; + if let Some(target) = target.clone() { + params.fts_target = Some(target); + } + let coordinate_rank = params + .fts_target() + .map(FtsTarget::coordinate_rank) + .unwrap_or(0); let partitions = reader .schema() .metadata @@ -1269,6 +1391,7 @@ impl InvertedIndex { frag_reuse_index_clone, &index_cache_for_part, token_set_format, + coordinate_rank, ) .await?, )) @@ -1292,8 +1415,21 @@ impl InvertedIndex { })) } Err(_) => { + if target.as_ref().is_some_and(FtsTarget::is_element_document) { + return Err(Error::not_supported( + "element-document FTS is not supported by the legacy inverted index format" + .to_string(), + )); + } // old index format - Self::load_legacy_index(store, frag_reuse_index, index_cache).await + let index = Self::load_legacy_index(store, frag_reuse_index, index_cache).await?; + if let Some(target) = target { + let mut index = (*index).clone(); + index.params.fts_target = Some(target); + Ok(Arc::new(index)) + } else { + Ok(index) + } } } } @@ -1735,6 +1871,7 @@ impl InvertedPartition { frag_reuse_index: Option>, index_cache: &LanceCache, token_set_format: TokenSetFormat, + coordinate_rank: usize, ) -> Result { let token_file = store.open_index_file(&token_file_path(id)).await?; let tokens = TokenSet::load(token_file, token_set_format).await?; @@ -1756,6 +1893,7 @@ impl InvertedPartition { frag_reuse_index, // V3 (256-doc block) partitions score with quantized doc lengths. inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, + coordinate_rank, )); Ok(Self { @@ -6110,6 +6248,9 @@ const fn build_dequantized_doc_lengths() -> [u32; 256] { pub struct DocSet { row_ids: Vec, num_tokens: Vec, + // One flat u32 column per list boundary. This avoids a Vec allocation per + // document while preserving the full logical document coordinate. + doc_indices: Vec>, // (row_id, doc_id) pairs sorted by row_id inv: Vec<(u64, u32)>, @@ -6127,6 +6268,7 @@ impl DeepSizeOf for DocSet { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { self.row_ids.deep_size_of_children(context) + self.num_tokens.deep_size_of_children(context) + + self.doc_indices.deep_size_of_children(context) + self.inv.deep_size_of_children(context) + self .norms @@ -6168,11 +6310,26 @@ impl DocSet { self.row_ids[doc_id as usize] } + pub fn doc_index(&self, doc_id: u32) -> Vec { + self.doc_indices + .iter() + .map(|coordinates| coordinates[doc_id as usize]) + .collect() + } + + pub fn coordinate(&self, doc_id: u32, rank: usize) -> u32 { + self.doc_indices[rank][doc_id as usize] + } + + pub fn coordinate_rank(&self) -> usize { + self.doc_indices.len() + } + /// Resolve a `row_id` to every `doc_id` it owns. /// - /// Modern indexes map each row to a single document. Older list indexes - /// may have indexed each list element as its own document, so a single - /// `row_id` can still own several `doc_id`s sharing that key in `inv`. + /// Row-document indexes map each row to a single document. Element-document + /// indexes (and older list indexes) can map one row to several documents, + /// so a single `row_id` may own multiple `doc_id`s sharing that key in `inv`. /// The prefilter path (`flat_search`) walks an allow-list of row_ids and /// must evaluate all legacy documents for that row. pub fn doc_ids(&self, row_id: u64) -> impl Iterator + '_ { @@ -6242,18 +6399,27 @@ impl DocSet { let row_id_col = UInt64Array::from_iter_values(self.row_ids.iter().cloned()); let num_tokens_col = UInt32Array::from_iter_values(self.num_tokens.iter().cloned()); - let schema = arrow_schema::Schema::new(vec![ + let mut fields = vec![ arrow_schema::Field::new(ROW_ID, DataType::UInt64, false), arrow_schema::Field::new(NUM_TOKEN_COL, DataType::UInt32, false), - ]); + ]; + let mut columns = vec![ + Arc::new(row_id_col) as ArrayRef, + Arc::new(num_tokens_col) as ArrayRef, + ]; + for (rank, coordinates) in self.doc_indices.iter().enumerate() { + fields.push(arrow_schema::Field::new( + doc_index_storage_column(rank), + DataType::UInt32, + false, + )); + columns.push( + Arc::new(UInt32Array::from_iter_values(coordinates.iter().copied())) as ArrayRef, + ); + } + let schema = arrow_schema::Schema::new(fields); - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(row_id_col) as ArrayRef, - Arc::new(num_tokens_col) as ArrayRef, - ], - )?; + let batch = RecordBatch::try_new(Arc::new(schema), columns)?; Ok(batch) } @@ -6265,7 +6431,21 @@ impl DocSet { let batch = reader.read_range(0..reader.num_rows(), None).await?; let row_id_col = batch[ROW_ID].as_primitive::(); let num_tokens_col = batch[NUM_TOKEN_COL].as_primitive::(); - Self::from_columns(row_id_col, num_tokens_col, is_legacy, frag_reuse_index) + let mut doc_indices = Vec::new(); + for rank in 0.. { + let column_name = doc_index_storage_column(rank); + let Some(column) = batch.column_by_name(&column_name) else { + break; + }; + doc_indices.push(column.as_primitive::()); + } + Self::from_columns_with_doc_indices( + row_id_col, + num_tokens_col, + &doc_indices, + is_legacy, + frag_reuse_index, + ) } /// Build a `DocSet` carrying only the per-doc `num_tokens` array; @@ -6279,6 +6459,7 @@ impl DocSet { Self { row_ids: Vec::new(), num_tokens, + doc_indices: Vec::new(), inv: Vec::new(), total_tokens, scoring_quantized: false, @@ -6296,6 +6477,34 @@ impl DocSet { is_legacy: bool, frag_reuse_index: Option>, ) -> Result { + Self::from_columns_with_doc_indices( + row_id_col, + num_tokens_col, + &[], + is_legacy, + frag_reuse_index, + ) + } + + pub fn from_columns_with_doc_indices( + row_id_col: &UInt64Array, + num_tokens_col: &arrow_array::UInt32Array, + doc_index_cols: &[&arrow_array::UInt32Array], + is_legacy: bool, + frag_reuse_index: Option>, + ) -> Result { + if doc_index_cols + .iter() + .any(|column| column.len() != row_id_col.len()) + { + return Err(Error::index( + "FTS document coordinate columns must have the same length as row ids".to_string(), + )); + } + let doc_indices = doc_index_cols + .iter() + .map(|column| column.values().to_vec()) + .collect::>(); // for legacy format, the row id is doc id; sorting keeps binary search viable if is_legacy { let (row_ids, num_tokens): (Vec<_>, Vec<_>) = row_id_col @@ -6316,6 +6525,7 @@ impl DocSet { return Ok(Self { row_ids, num_tokens, + doc_indices, inv: Vec::new(), total_tokens, scoring_quantized: false, @@ -6359,6 +6569,7 @@ impl DocSet { return Ok(Self { row_ids, num_tokens, + doc_indices, inv, total_tokens, scoring_quantized: false, @@ -6380,6 +6591,7 @@ impl DocSet { Ok(Self { row_ids, num_tokens, + doc_indices, inv, total_tokens, scoring_quantized: false, @@ -6394,6 +6606,11 @@ impl DocSet { let len = self.len(); let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, Vec::with_capacity(len)); + let doc_indices = std::mem::take(&mut self.doc_indices); + self.doc_indices = doc_indices + .iter() + .map(|_| Vec::with_capacity(len)) + .collect(); self.invalidate_norms(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { @@ -6401,6 +6618,11 @@ impl DocSet { Some(Some(new_row_id)) => { self.row_ids.push(new_row_id); self.num_tokens.push(num_token); + for (new_coordinates, old_coordinates) in + self.doc_indices.iter_mut().zip(&doc_indices) + { + new_coordinates.push(old_coordinates[doc_id]); + } self.total_tokens += num_token as u64; } Some(None) => { @@ -6409,6 +6631,11 @@ impl DocSet { None => { self.row_ids.push(row_id); self.num_tokens.push(num_token); + for (new_coordinates, old_coordinates) in + self.doc_indices.iter_mut().zip(&doc_indices) + { + new_coordinates.push(old_coordinates[doc_id]); + } self.total_tokens += num_token as u64; } } @@ -6474,6 +6701,32 @@ impl DocSet { self.row_ids.len() as u32 - 1 } + pub fn append_with_doc_index( + &mut self, + row_id: u64, + num_tokens: u32, + doc_index: &[u32], + ) -> Result { + if self.row_ids.is_empty() && self.doc_indices.is_empty() { + self.doc_indices = (0..doc_index.len()).map(|_| Vec::new()).collect(); + } + if self.doc_indices.len() != doc_index.len() { + return Err(Error::index(format!( + "all documents in an FTS partition must have the same coordinate rank: expected {}, got {}", + self.doc_indices.len(), + doc_index.len() + ))); + } + self.row_ids.push(row_id); + self.num_tokens.push(num_tokens); + for (coordinates, value) in self.doc_indices.iter_mut().zip(doc_index) { + coordinates.push(*value); + } + self.total_tokens += num_tokens as u64; + self.invalidate_norms(); + Ok(self.row_ids.len() as u32 - 1) + } + // Drop the baked norm slab after a mutation; it re-bakes on the next // scoring use. fn invalidate_norms(&mut self) { @@ -6485,10 +6738,19 @@ impl DocSet { pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.capacity() * std::mem::size_of::() + + self + .doc_indices + .iter() + .map(|coordinates| coordinates.capacity() * std::mem::size_of::()) + .sum::() + self.inv.capacity() * std::mem::size_of::<(u64, u32)>() } } +pub fn doc_index_storage_column(rank: usize) -> String { + format!("{DOC_INDEX_STORAGE_PREFIX}{rank}") +} + pub fn flat_full_text_search( batches: &[&RecordBatch], doc_col: &str, @@ -6589,9 +6851,8 @@ fn do_flat_full_text_search_list( Ok(results) } -const FLAT_ROW_ID_COL_IDX: usize = 0; -const FLAT_ALL_TOKENS_COL_IDX: usize = 1; -const FLAT_QUERY_TOKEN_COUNTS_COL_IDX: usize = 2; +const FLAT_ALL_TOKENS_COL: &str = "all_tokens"; +const FLAT_QUERY_TOKEN_COUNTS_COL: &str = "query_token_counts"; /// If we accumulate this many bytes we warn the user they probably want to use an FTS index instead. const BYTES_ACCUMULATED_WARNING_THRESHOLD: u64 = 1024 * 1024 * 1024; // 1GB @@ -6615,19 +6876,28 @@ async fn tokenize_and_count( query_tokens: Arc, doc_col_idx: usize, elapsed_compute: Option