From 393974fac037fbdb8293412c55f1879aceed382c Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sun, 21 Jun 2026 21:12:19 +0000 Subject: [PATCH 1/5] feat(index): support flattened JSON sub-doc indexing Add explicit JsonTokenizerMode values, with SingleDocument for existing JSON indexes and FlattenedSubDocs for new JSON indexes. Flatten JSON arrays into multiple internal inverted-index sub-docs while mapping each sub-doc back to the original row id. Normalize bracketed JSON query paths into value tokens plus constraint tokens, and deduplicate flattened JSON search results by row id. Add persisted disable_cross_array_unnest index metadata, defaulting to false. When set, flattened JSON tokenization indexes sibling arrays independently instead of producing their Cartesian product, matching Pinot's memory-saving DisableCrossArrayUnnest behavior. Expose disable_cross_array_unnest through Rust, Python, and Java inverted-index params. Test Plan: PASS: cargo fmt --all --check PASS: cargo test -p lance-index scalar::inverted::tokenizer::document_tokenizer::tests --lib PASS: cargo test -p lance-index scalar::inverted::tokenizer::tests --lib PASS: cargo test -p lance test_json_inverted_ --lib PASS: cargo test -p lance test_auto_infer_lance_tokenizer --lib PASS: cargo check -p lance-index --tests PASS: cargo check -p lance --tests PASS: cargo clippy --all --tests --benches -- -D warnings PASS: PATH=/home/user/.cargo/bin:$PATH cargo fmt --manifest-path ./lance-jni/Cargo.toml --all --check (from java/) PASS: ./mvnw spotless:check (from java/) PASS: PATH=/home/user/.cargo/bin:$PATH ./mvnw test (from java/) PASS: PATH=/home/user/.local/bin:$PATH uv run pytest python/tests/test_scalar_index.py::test_json_inverted_match_query (from python/) PASS: PATH=/home/user/.local/bin:/home/user/.cargo/bin:$PATH uv run make lint (from python/) NOTE: PATH=/home/user/.local/bin:$PATH make install (from python/) built pylance and installed deps, then failed at pre-commit install because core.hooksPath is set to /etc/git-hooks. --- .../index/scalar/InvertedIndexParams.java | 20 + protos/index_old.proto | 6 + python/python/lance/dataset.py | 6 + python/python/tests/test_scalar_index.py | 5 +- python/src/dataset.rs | 18 + rust/lance-index/src/scalar/inverted.rs | 2 +- .../src/scalar/inverted/builder.rs | 202 +++---- .../src/scalar/inverted/compound.rs | 13 +- rust/lance-index/src/scalar/inverted/index.rs | 1 + .../src/scalar/inverted/index/flat_search.rs | 264 ++++---- .../src/scalar/inverted/index/search.rs | 61 +- rust/lance-index/src/scalar/inverted/query.rs | 16 +- .../src/scalar/inverted/tokenizer.rs | 64 +- .../inverted/tokenizer/document_tokenizer.rs | 567 +++++++++++++++++- rust/lance/src/dataset/tests/dataset_index.rs | 84 ++- rust/lance/src/io/exec/fts.rs | 50 +- 16 files changed, 1123 insertions(+), 256 deletions(-) diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index 6597fdc0841..852ba18ffcd 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -57,6 +57,7 @@ public static final class Builder { private Integer maxNgramLength; private Boolean prefixOnly; private Integer blockSize = 128; + private Boolean disableCrossArrayUnnest; private Boolean splitIdentifiers; private Boolean splitOnNumerics; private Boolean preserveOriginal; @@ -304,6 +305,22 @@ public Builder blockSize(int blockSize) { return this; } + /** + * Configure whether flattened JSON tokenization avoids cross-array unnesting. + * + *

When true, sibling arrays are indexed independently instead of producing their Cartesian + * product. This can reduce index build memory for JSON records with multiple arrays but can + * sacrifice result accuracy for queries that constrain values across those arrays. The default + * is false. + * + * @param disableCrossArrayUnnest whether to avoid cross-array unnesting + * @return this builder + */ + public Builder disableCrossArrayUnnest(boolean disableCrossArrayUnnest) { + this.disableCrossArrayUnnest = disableCrossArrayUnnest; + return this; + } + /** * Configure whether code identifiers are split into subwords. * @@ -501,6 +518,9 @@ public ScalarIndexParams build() { if (blockSize != null) { params.put("block_size", blockSize); } + if (disableCrossArrayUnnest != null) { + params.put("disable_cross_array_unnest", disableCrossArrayUnnest); + } if (splitIdentifiers != null) { params.put("split_identifiers", splitIdentifiers); } diff --git a/protos/index_old.proto b/protos/index_old.proto index 01be98faeb0..0f7a84fe3cd 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -115,4 +115,10 @@ message InvertedIndexDetails { * which identifies the overall inverted-index layout. */ optional uint32 posting_format_version = 15; + // JSON document tokenization mode. Absent means SingleDocument JSON tokenization, + // which is how indexes written before flattened JSON sub-docs are interpreted. + optional string json_tokenizer_mode = 16; + // If true, avoid cross-array unnesting during flattened JSON tokenization. + // The default false value preserves exact Cartesian-product semantics. + bool disable_cross_array_unnest = 17; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 6ae17bb1c00..9f4dc3cb093 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3720,6 +3720,12 @@ def create_scalar_index( ``[1, num_compute_cpus]``. If unset, Lance uses ``num_compute_cpus`` workers unless ``LANCE_FTS_NUM_SHARDS`` is set. This parameter is only used for the current build and is not persisted with the index. + disable_cross_array_unnest: bool, default False + This is for the ``INVERTED`` index on JSON columns. If True, flattened + JSON tokenization indexes sibling arrays independently instead of + producing their Cartesian product. This reduces index build memory for + records with multiple arrays but can sacrifice result accuracy for + queries that constrain values across those arrays. base_tokenizer: str, default "simple" This is for the ``INVERTED`` index. The base tokenizer to use. The value can be: diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 79762a9539f..6d7bfba1747 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -6032,7 +6032,10 @@ def test_json_inverted_match_query(tmp_path): stem=True, lower_case=True, remove_stop_words=True, + disable_cross_array_unnest=True, ) + details = dataset.describe_indices()[0].details + assert details["disable_cross_array_unnest"] is True # Test match query with token exceeding max_token_length results = dataset.to_table( @@ -6048,7 +6051,7 @@ def test_json_inverted_match_query(tmp_path): # Test language match results = dataset.to_table( - full_text_query=MatchQuery("Language,str,english", "json_col") + full_text_query=MatchQuery("Language[*],str,english", "json_col") ) assert results.num_rows == 3 diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 547ba91d16b..ae75f92734d 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2682,6 +2682,24 @@ impl Dataset { .block_size(block_size.extract()?) .map_err(|e| PyValueError::new_err(e.to_string()))?; } + if let Some(split_identifiers) = kwargs.get_item("split_identifiers")? { + params = params.split_identifiers(split_identifiers.extract()?); + } + if let Some(split_on_numerics) = kwargs.get_item("split_on_numerics")? { + params = params.split_on_numerics(split_on_numerics.extract()?); + } + if let Some(preserve_original) = kwargs.get_item("preserve_original")? { + params = params.preserve_original(preserve_original.extract()?); + } + if let Some(index_operators) = kwargs.get_item("index_operators")? { + params = params.index_operators(index_operators.extract()?); + } + if let Some(disable_cross_array_unnest) = + kwargs.get_item("disable_cross_array_unnest")? + { + params = params + .disable_cross_array_unnest(disable_cross_array_unnest.extract()?); + } if let Some(memory_limit) = kwargs.get_item("memory_limit")? { params = params.memory_limit_mb(memory_limit.extract()?); } diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index c3070d4d61c..b26824e6c27 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -428,11 +428,11 @@ impl InvertedIndexPlugin { params.validate_format_version()?; let format_version = params.resolved_format_version(); let is_element_document = params.get_document_granularity().is_list_element(); - let details = pbold::InvertedIndexDetails::try_from(¶ms)?; let mut inverted_index = InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask) .with_progress(progress); let files = inverted_index.update(data, index_store, None).await?; + let details = pbold::InvertedIndexDetails::try_from(inverted_index.params())?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&details).unwrap(), index_version: if is_element_document { diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 10e613678c2..7681843b7b9 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use super::{InvertedIndexParams, index::*}; -use crate::scalar::inverted::document_tokenizer::DocType; +use crate::scalar::inverted::document_tokenizer::{DocType, JsonTokenizerMode}; use crate::scalar::inverted::json::JsonTextStream; use crate::scalar::inverted::tokenizer::LEGACY_BLOCK_SIZE; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; @@ -183,13 +183,19 @@ impl InvertedIndexBuilder { /// Constructed as `(fragment_id as u64) << 32`. /// When provided, ensures that generated IDs belong to the specified fragment. pub fn from_existing_index( - params: InvertedIndexParams, + mut params: InvertedIndexParams, store: Option>, partitions: Vec, token_set_format: TokenSetFormat, fragment_mask: Option, deleted_fragments: RoaringBitmap, ) -> Self { + if (store.is_some() || !partitions.is_empty()) + && params.lance_tokenizer.as_deref() == Some("json") + && params.json_tokenizer_mode.is_none() + { + params.json_tokenizer_mode = Some(JsonTokenizerMode::SingleDocument); + } let format_version = params.resolved_format_version(); Self { params, @@ -205,6 +211,17 @@ impl InvertedIndexBuilder { } } + fn configure_json_tokenizer_mode_for_new_data(&mut self, doc_type: DocType) { + if self.params.lance_tokenizer.is_none() { + self.params.lance_tokenizer = Some(doc_type.as_ref().to_string()); + } + if self.params.lance_tokenizer.as_deref() == Some("json") + && self.params.json_tokenizer_mode.is_none() + { + self.params.json_tokenizer_mode = Some(JsonTokenizerMode::FlattenedSubDocs); + } + } + pub fn with_posting_tail_codec(mut self, posting_tail_codec: PostingTailCodec) -> Self { self.format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( posting_tail_codec, @@ -231,6 +248,10 @@ impl InvertedIndexBuilder { self } + pub(crate) fn params(&self) -> &InvertedIndexParams { + &self.params + } + pub async fn update( &mut self, new_data: SendableRecordBatchStream, @@ -242,12 +263,9 @@ impl InvertedIndexBuilder { let doc_col = schema.field(0).name(); // infer lance_tokenizer based on document type - if self.params.lance_tokenizer.is_none() { - let schema = new_data.schema(); - let field = schema.column_with_name(doc_col).expect_ok()?.1; - let doc_type = DocType::try_from(field)?; - self.params.lance_tokenizer = Some(doc_type.as_ref().to_string()); - } + let field = schema.column_with_name(doc_col).expect_ok()?.1; + let doc_type = DocType::try_from(field)?; + self.configure_json_tokenizer_mode_for_new_data(doc_type); let new_data = document_input(new_data, doc_col)?; @@ -276,11 +294,9 @@ impl InvertedIndexBuilder { let schema = new_data.schema(); let doc_col = schema.field(0).name(); - if self.params.lance_tokenizer.is_none() { - let field = schema.column_with_name(doc_col).expect_ok()?.1; - let doc_type = DocType::try_from(field)?; - self.params.lance_tokenizer = Some(doc_type.as_ref().to_string()); - } + let field = schema.column_with_name(doc_col).expect_ok()?.1; + let doc_type = DocType::try_from(field)?; + self.configure_json_tokenizer_mode_for_new_data(doc_type); let mut files = self .merge_existing_segments(dest_store, old_segments, old_data_filter.as_ref()) @@ -1576,106 +1592,90 @@ impl IndexWorker { document: DocumentSource<'_>, doc_index: &[u32], ) -> Result<()> { + let doc = match document { + DocumentSource::Text(doc) => doc.to_string(), + DocumentSource::StringList(elements) => Self::materialize_string_list(elements), + }; + self.total_doc_length += doc.len(); let with_position = self.has_position(); + let sub_docs = self.tokenizer.token_streams_for_doc(&doc)?; + for tokens in sub_docs { + self.process_tokenized_doc(row_id, tokens, with_position, doc_index) + .await?; + } + + Ok(()) + } + + async fn process_tokenized_doc( + &mut self, + row_id: u64, + mut tokens: Vec, + with_position: bool, + doc_index: &[u32], + ) -> Result<()> { let builder_was_empty = self.builder.docs.is_empty(); let old_temporary_memory_size = self.temporary_memory_size(); let old_token_memory_size = self.builder.tokens.memory_size() as u64; let doc_id = self.builder.docs.len() as u32; let mut token_num: u32 = 0; - let mut doc_length_bytes = 0usize; let mut posting_memory_delta = 0i64; if with_position { - { - if self.token_ids.capacity() < self.last_token_count { - self.token_ids - .reserve(self.last_token_count - self.token_ids.capacity()); + if self.token_ids.capacity() < self.last_token_count { + self.token_ids + .reserve(self.last_token_count - self.token_ids.capacity()); + } + self.token_ids.clear(); + let builder = &mut self.builder; + let token_ids = &mut self.token_ids; + let memory_size = &mut self.memory_size; + let posting_tail_codec = builder.posting_tail_codec; + let block_size = builder.block_size; + + for token in &mut tokens { + let position = Self::checked_token_position(row_id, token.position)?; + let token_id = builder.tokens.get_or_add(&token.text); + if token_id as usize == builder.posting_lists.len() { + let old_posting_lists_overhead_size = (builder.posting_lists.capacity() + * std::mem::size_of::()) + as u64; + builder.posting_lists.push( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + true, + posting_tail_codec, + block_size, + ), + ); + let new_posting_lists_overhead_size = (builder.posting_lists.capacity() + * std::mem::size_of::()) + as u64; + Self::adjust_tracked_value( + memory_size, + old_posting_lists_overhead_size, + new_posting_lists_overhead_size, + ); } - self.token_ids.clear(); - let tokenizer = &mut self.tokenizer; - let builder = &mut self.builder; - let token_ids = &mut self.token_ids; - let memory_size = &mut self.memory_size; - let posting_tail_codec = builder.posting_tail_codec; - - let block_size = builder.block_size; - let mut process_text = |text: &str| -> Result<()> { - doc_length_bytes += text.len(); - let mut token_stream = tokenizer.token_stream_for_doc(text); - while token_stream.advance() { - let token = token_stream.token(); - let position = Self::checked_token_position(row_id, token.position)?; - let token_id = builder.tokens.get_or_add(&token.text); - if token_id as usize == builder.posting_lists.len() { - let old_posting_lists_overhead_size = (builder.posting_lists.capacity() - * std::mem::size_of::()) - as u64; - builder.posting_lists.push( - PostingListBuilder::new_with_posting_tail_codec_and_block_size( - true, - posting_tail_codec, - block_size, - ), - ); - let new_posting_lists_overhead_size = (builder.posting_lists.capacity() - * std::mem::size_of::()) - as u64; - Self::adjust_tracked_value( - memory_size, - old_posting_lists_overhead_size, - new_posting_lists_overhead_size, - ); - } - let posting_list = &mut builder.posting_lists[token_id as usize]; - let old_posting_memory_size = posting_list.size(); - if posting_list.add_occurrence(doc_id, position)? { - token_ids.push(token_id); - } - let new_posting_memory_size = posting_list.size(); - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; - token_num += 1; - } - Ok(()) - }; - - match document { - DocumentSource::Text(doc) => { - process_text(doc)?; - } - DocumentSource::StringList(elements) => { - let doc = Self::materialize_string_list(elements); - process_text(&doc)?; - } + let posting_list = &mut builder.posting_lists[token_id as usize]; + let old_posting_memory_size = posting_list.size(); + if posting_list.add_occurrence(doc_id, position)? { + token_ids.push(token_id); } + let new_posting_memory_size = posting_list.size(); + posting_memory_delta += + new_posting_memory_size as i64 - old_posting_memory_size as i64; + token_num += 1; } } else { - { - if self.token_ids.capacity() < self.last_token_count { - self.token_ids - .reserve(self.last_token_count - self.token_ids.capacity()); - } - self.token_ids.clear(); - - let tokenizer = &mut self.tokenizer; - let builder = &mut self.builder; - let token_ids = &mut self.token_ids; - let mut process_text = |text: &str| { - doc_length_bytes += text.len(); - let mut token_stream = tokenizer.token_stream_for_doc(text); - while token_stream.advance() { - let token_id = builder.tokens.get_or_add(&token_stream.token().text); - token_ids.push(token_id); - token_num += 1; - } - }; + if self.token_ids.capacity() < self.last_token_count { + self.token_ids + .reserve(self.last_token_count - self.token_ids.capacity()); + } + self.token_ids.clear(); - match document { - DocumentSource::Text(doc) => process_text(doc), - DocumentSource::StringList(elements) => { - let doc = Self::materialize_string_list(elements); - process_text(&doc); - } - } + for token in tokens { + let token_id = self.builder.tokens.get_or_add(&token.text); + self.token_ids.push(token_id); + token_num += 1; } } self.adjust_tracked_memory_size( @@ -1728,7 +1728,6 @@ impl IndexWorker { old_doc_memory_size, self.builder.docs.memory_size() as u64, ); - self.total_doc_length += doc_length_bytes; if with_position { for &token_id in &self.token_ids { @@ -1793,7 +1792,6 @@ impl IndexWorker { { self.flush().await?; } - Ok(()) } diff --git a/rust/lance-index/src/scalar/inverted/compound.rs b/rust/lance-index/src/scalar/inverted/compound.rs index 7bbba53ee52..b0c4b485218 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -17,7 +17,7 @@ use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; use super::{ InvertedIndex, PreparedBm25Query, - document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}, + document_tokenizer::{DocType, JsonTokenizer, JsonTokenizerMode, LanceTokenizer}, documents::{ CachedRowAddressOrder, DocId, DocLengths, DocVisibility, OrderedRowAddressProjection, PartitionDocuments, ResidentAddressProjection, RowAddressProjectionOrderError, @@ -3791,9 +3791,16 @@ pub(super) fn tokenize_leaf( && matches!(params.fuzziness, Some(distance) if distance > 0); let mut tokenizer = if is_explicit_fuzzy_match { let analyzer = TextAnalyzer::from(SimpleTokenizer::default()); - match index.tokenizer().doc_type() { + let index_tokenizer = index.tokenizer(); + match index_tokenizer.doc_type() { DocType::Text => Box::new(TextTokenizer::new(analyzer)) as Box, - DocType::Json => Box::new(JsonTokenizer::new(analyzer)) as Box, + DocType::Json => Box::new(JsonTokenizer::new( + analyzer, + index_tokenizer + .json_tokenizer_mode() + .unwrap_or(JsonTokenizerMode::SingleDocument), + index_tokenizer.disable_cross_array_unnest(), + )) as Box, } } else { index.tokenizer() diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 660339fefd3..fcd3bb2a285 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -65,6 +65,7 @@ use super::documents::{ use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder}; use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder}; use super::iter::PostingListIterator; +use super::tokenizer::document_tokenizer::JsonTokenizerMode; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; use super::{DocumentGranularity, InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ diff --git a/rust/lance-index/src/scalar/inverted/index/flat_search.rs b/rust/lance-index/src/scalar/inverted/index/flat_search.rs index 3417214a7f5..97ba79a0ee4 100644 --- a/rust/lance-index/src/scalar/inverted/index/flat_search.rs +++ b/rust/lance-index/src/scalar/inverted/index/flat_search.rs @@ -299,87 +299,99 @@ pub(super) async fn tokenize_and_count( .collect::>(); let mut phrase_matches = phrase_slop .map(|_| BooleanBuilder::with_capacity(batch.num_rows())); - let mut count_text = |doc: &str, - temp_query_token_counts: &mut Vec| - -> DataFusionResult<(u64, bool)> { - for positions in &mut temp_query_positions { - positions.clear(); - } - let mut stream = tokenizer.token_stream_for_doc(doc); - let mut all_tokens = 0; - while let Some(token) = stream.next() { - all_tokens += 1; - if let Some(token_indices) = query_token_indices.get(&token.text) { - for token_index in token_indices { - temp_query_token_counts[*token_index] += 1; - if phrase_slop.is_some() { - temp_query_positions[*token_index].push( - u32::try_from(token.position).map_err(|_| { - datafusion_common::DataFusionError::Execution(format!( - "flat FTS token position exceeds u32: {}", - token.position - )) - })?, - ); - } - } - } - } - let matches = phrase_slop.is_none_or(|slop| { - phrase_matches_positions( - query_tokens.as_ref(), - &temp_query_positions, - slop, - ) - }); - Ok((all_tokens, matches)) - }; - let mut append_counts = |row_index: usize, - row_id: u64, - all_tokens: u64, - temp_query_token_counts: &[u64], - phrase_match: bool| - -> DataFusionResult<()> { - row_ids.append_value(row_id); - for (builder, input) in - doc_indices.iter_mut().zip(input_doc_indices.iter()) + macro_rules! append_counts { + ($row_index:expr, $row_id:expr, $all_tokens:expr, $counts:expr, $phrase_match:expr $(,)?) => {{ + row_ids.append_value($row_id); + for (builder, input) in doc_indices.iter_mut().zip(input_doc_indices.iter()) { - builder.append_value(input.value(row_index)); + builder.append_value(input.value($row_index)); } - all_token_counts.append_value(all_tokens); - for count in temp_query_token_counts.iter().copied() { + all_token_counts.append_value($all_tokens); + for count in $counts.iter().copied() { query_token_counts.values().append_value(count); } query_token_counts.append(true); if let Some(builder) = phrase_matches.as_mut() { - builder.append_value(phrase_match); + builder.append_value($phrase_match); } - Ok(()) - }; - match batch.column(doc_col_idx).data_type() { - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { - let doc_iter = iter_str_array(batch.column(doc_col_idx)); - for (row_index, (doc, row_id)) in - doc_iter.zip(row_id_array.values().iter()).enumerate() - { + }}; + } + let mut append_doc = + |row_index: usize, row_id: u64, doc: Option<&str>| -> DataFusionResult<()> { + let Some(doc) = doc else { + if coordinate_rank > 0 { + temp_query_token_counts.clear(); + temp_query_token_counts + .extend(std::iter::repeat_n(0, query_tokens.len())); + append_counts!( + row_index, + row_id, + 0, + &temp_query_token_counts, + false, + ); + } + return Ok(()); + }; + let sub_docs = tokenizer.token_streams_for_doc(doc).map_err(|err| { + datafusion_common::DataFusionError::Execution(err.to_string()) + })?; + for tokens in sub_docs { temp_query_token_counts.clear(); temp_query_token_counts .extend(std::iter::repeat_n(0, query_tokens.len())); + let mut all_tokens = 0; - let (all_tokens, phrase_match) = match doc { - Some(doc) => count_text(doc, &mut temp_query_token_counts)?, - None => (0, false), - }; + for positions in &mut temp_query_positions { + positions.clear(); + } + for token in tokens { + all_tokens += 1; + if let Some(token_indices) = query_token_indices.get(&token.text) { + for token_index in token_indices { + temp_query_token_counts[*token_index] += 1; + if phrase_slop.is_some() { + temp_query_positions[*token_index].push( + u32::try_from(token.position).map_err(|_| { + datafusion_common::DataFusionError::Execution( + format!( + "flat FTS token position exceeds u32: {}", + token.position + ), + ) + })?, + ); + } + } + } + } + let phrase_match = phrase_slop.is_none_or(|slop| { + phrase_matches_positions( + query_tokens.as_ref(), + &temp_query_positions, + slop, + ) + }); if coordinate_rank > 0 || all_tokens > 0 { - append_counts( + append_counts!( row_index, - *row_id, + row_id, all_tokens, &temp_query_token_counts, phrase_match, - )?; + ); } } + Ok(()) + }; + match batch.column(doc_col_idx).data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + let doc_iter = iter_str_array(batch.column(doc_col_idx)); + for (row_index, (doc, row_id)) in + doc_iter.zip(row_id_array.values().iter()).enumerate() + { + append_doc(row_index, *row_id, doc)?; + } } DataType::List(_) => { if coordinate_rank != 0 { @@ -390,14 +402,10 @@ pub(super) async fn tokenize_and_count( ), ); } - tokenize_and_count_list::( + append_list_docs::( batch.column(doc_col_idx), row_id_array, - &mut count_text, - &mut append_counts, - &mut temp_query_token_counts, - query_tokens.len(), - phrase_slop.is_some(), + &mut append_doc, )?; } DataType::LargeList(_) => { @@ -409,14 +417,10 @@ pub(super) async fn tokenize_and_count( ), ); } - tokenize_and_count_list::( + append_list_docs::( batch.column(doc_col_idx), row_id_array, - &mut count_text, - &mut append_counts, - &mut temp_query_token_counts, - query_tokens.len(), - phrase_slop.is_some(), + &mut append_doc, )?; } data_type => { @@ -465,14 +469,10 @@ pub(super) async fn tokenize_and_count( )?) } -pub(super) fn tokenize_and_count_list( +pub(super) fn append_list_docs( doc_col: &ArrayRef, row_id_array: &arrow_array::PrimitiveArray, - count_text: &mut impl FnMut(&str, &mut Vec) -> DataFusionResult<(u64, bool)>, - append_counts: &mut impl FnMut(usize, u64, u64, &[u64], bool) -> DataFusionResult<()>, - temp_query_token_counts: &mut Vec, - query_tokens_len: usize, - match_phrase: bool, + append_doc: &mut impl FnMut(usize, u64, Option<&str>) -> DataFusionResult<()>, ) -> DataFusionResult<()> { let doc_array = doc_col.as_list::(); match doc_array.value_type() { @@ -486,35 +486,20 @@ pub(super) fn tokenize_and_count_list( } for i in 0..row_id_array.len() { - temp_query_token_counts.clear(); - temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens_len)); - let mut all_tokens = 0; - let mut phrase_match = false; - if !doc_array.is_null(i) { - let elements = doc_array.value(i); - if match_phrase { - let mut document = String::new(); - for element in iter_str_array(elements.as_ref()).flatten() { - if !document.is_empty() { - document.push(' '); - } - document.push_str(element); - } - (all_tokens, phrase_match) = count_text(&document, temp_query_token_counts)?; - } else { - for element in iter_str_array(elements.as_ref()).flatten() { - all_tokens += count_text(element, temp_query_token_counts)?.0; - } + if doc_array.is_null(i) { + continue; + } + + let elements = doc_array.value(i); + let mut doc = String::new(); + for element in iter_str_array(elements.as_ref()).flatten() { + if !doc.is_empty() { + doc.push(' '); } + doc.push_str(element); } - if all_tokens > 0 { - append_counts( - i, - row_id_array.value(i), - all_tokens, - temp_query_token_counts, - phrase_match, - )?; + if !doc.is_empty() { + append_doc(i, row_id_array.value(i), Some(&doc))?; } } @@ -587,6 +572,49 @@ pub(super) fn initialize_scorer( MemBM25Scorer::new(total_tokens, num_docs, token_counts_map) } +fn deduplicate_scored_rows( + row_ids: Vec, + scores: Vec, + limit: usize, +) -> (Vec, Vec) { + let mut scores_by_row_id: HashMap = HashMap::with_capacity(row_ids.len()); + for (row_id, score) in row_ids.into_iter().zip(scores) { + scores_by_row_id + .entry(row_id) + .and_modify(|existing| { + if score > *existing { + *existing = score; + } + }) + .or_insert(score); + } + + let mut scored_rows = scores_by_row_id.into_iter().collect::>(); + scored_rows.sort_unstable_by(|(left_row_id, left_score), (right_row_id, right_score)| { + right_score + .total_cmp(left_score) + .then_with(|| left_row_id.cmp(right_row_id)) + }); + scored_rows.truncate(scored_rows.len().min(limit)); + scored_rows.into_iter().unzip() +} + +fn deduplicate_fts_batch(batch: RecordBatch, limit: usize) -> Result { + let row_ids = batch[ROW_ID].as_primitive::().values().to_vec(); + let scores = batch[SCORE_COL] + .as_primitive::() + .values() + .to_vec(); + let (row_ids, scores) = deduplicate_scored_rows(row_ids, scores, limit); + Ok(RecordBatch::try_new( + FTS_SCHEMA.clone(), + vec![ + Arc::new(UInt64Array::from(row_ids)) as ArrayRef, + Arc::new(Float32Array::from(scores)) as ArrayRef, + ], + )?) +} + pub(super) fn flat_bm25_score( query_tokens: &Tokens, counted_input: &RecordBatch, @@ -595,6 +623,7 @@ pub(super) fn flat_bm25_score( operator: Operator, boost: f32, phrase_slop: Option, + require_all_query_tokens: bool, ) -> Result { let mut row_ids_builder = UInt64Builder::with_capacity(counted_input.num_rows()); let mut scores_builder = Float32Builder::with_capacity(counted_input.num_rows()); @@ -680,12 +709,16 @@ pub(super) fn flat_bm25_score( } let doc_norm = K1 * (1.0 - B + B * num_tokens_in_doc as f32 / scorer.avg_doc_length()); let mut score = 0.0; + let mut has_all_query_tokens = true; for (token, freq) in query_tokens.into_iter().zip(query_token_counts) { let freq = freq as f32; + if freq == 0.0 { + has_all_query_tokens = false; + } let idf = idf(scorer.num_docs_containing_token(token), scorer.num_docs()); score += idf * (freq * (K1 + 1.0) / (freq + doc_norm)); } - if score > 0.0 { + if score > 0.0 && (!require_all_query_tokens || has_all_query_tokens) { row_ids_builder.append_value(row_id); if let Some(builder) = doc_indices_builder.as_mut() { for input_doc_index in &input_doc_indices { @@ -876,6 +909,13 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( // Pre-await synchronous work: query tokenization + chunk-stream setup. let pre_await_start = std::time::Instant::now(); let query_tokens = Arc::new(collect_query_tokens(&query, &mut tokenizer)); + let should_deduplicate_rows = + tokenizer.json_tokenizer_mode() == Some(JsonTokenizerMode::FlattenedSubDocs); + let require_all_query_tokens = should_deduplicate_rows + && query_tokens + .as_ref() + .into_iter() + .any(|token| token.contains("$idx,number,")); // A query that tokenizes to no terms (e.g. only stop words) has no // searchable content and matches nothing. Return early rather than @@ -937,7 +977,7 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( // All post-await work is synchronous; time the scorer + score + slicing loop together. let post_await_start = std::time::Instant::now(); let scorer = initialize_scorer(base_scorer.as_ref(), query_tokens.as_ref(), &counted_input); - let scores = flat_bm25_score( + let mut scores = flat_bm25_score( query_tokens.as_ref(), &counted_input, &scorer, @@ -945,7 +985,11 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( operator, boost, phrase_slop, + require_all_query_tokens, )?; + if should_deduplicate_rows { + scores = deduplicate_fts_batch(scores, usize::MAX)?; + } // Finally we emit batches according to the target batch size let num_out_batches = scores.num_rows().div_ceil(target_batch_size); diff --git a/rust/lance-index/src/scalar/inverted/index/search.rs b/rust/lance-index/src/scalar/inverted/index/search.rs index 78a533d7acb..6149d2b5d7c 100644 --- a/rust/lance-index/src/scalar/inverted/index/search.rs +++ b/rust/lance-index/src/scalar/inverted/index/search.rs @@ -587,22 +587,36 @@ impl InvertedIndex { || *LANCE_FTS_REUSE_PREPARED_SCORER_ENABLED, ); - let limit = params.limit.unwrap_or(usize::MAX); - if limit == 0 { + let requested_limit = params.limit.unwrap_or(usize::MAX); + if requested_limit == 0 { return Ok(Vec::new()); } + let should_deduplicate_rows = + self.params.json_tokenizer_mode == Some(JsonTokenizerMode::FlattenedSubDocs); + let search_limit = if should_deduplicate_rows { + usize::MAX + } else { + requested_limit + }; + let search_params = if should_deduplicate_rows { + let mut params = params.as_ref().clone(); + params.limit = None; + Arc::new(params) + } else { + params.clone() + }; let mask = prefilter.mask(); - if self.is_legacy() { + let documents = if self.is_legacy() { let (row_ids, scores) = self .bm25_search_legacy( tokens, - params, + search_params, operator, mask, metrics, scorer, impact_scorer, - limit, + search_limit, ) .await?; Ok(row_ids @@ -613,16 +627,21 @@ impl InvertedIndex { } else { self.bm25_search_modern(ModernSearchRequest { tokens, - params, + params: search_params, operator, mask, metrics, scorer, impact_scorer, - limit, + limit: search_limit, initial_score_floor, }) .await + }?; + if should_deduplicate_rows { + Ok(deduplicate_scored_documents(documents, requested_limit)) + } else { + Ok(documents) } } @@ -1168,6 +1187,34 @@ impl InvertedIndex { } } +fn deduplicate_scored_documents(documents: Vec, limit: usize) -> Vec { + let mut scores_by_row_id: HashMap = HashMap::with_capacity(documents.len()); + for document in documents { + let score = document.score.0; + scores_by_row_id + .entry(document.row_id) + .and_modify(|existing| { + if score > *existing { + *existing = score; + } + }) + .or_insert(score); + } + + let mut scored_documents = scores_by_row_id + .into_iter() + .map(|(row_id, score)| ScoredDoc::new(row_id, score)) + .collect::>(); + scored_documents.sort_unstable_by(|left, right| { + right + .score + .cmp(&left.score) + .then_with(|| left.row_id.cmp(&right.row_id)) + }); + scored_documents.truncate(scored_documents.len().min(limit)); + scored_documents +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/lance-index/src/scalar/inverted/query.rs b/rust/lance-index/src/scalar/inverted/query.rs index 9952c36cc55..e0a19056856 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -904,13 +904,19 @@ pub fn has_query_token( tokenizer: &mut Box, query_tokens: &Tokens, ) -> bool { - let mut stream = tokenizer.token_stream_for_doc(text); - while let Some(token) = stream.next() { - if query_tokens.contains(&token.text) { - return true; + match tokenizer.token_streams_for_doc(text) { + Ok(sub_docs) => { + for tokens in sub_docs { + for token in tokens { + if query_tokens.contains(&token.text) { + return true; + } + } + } + false } + Err(_) => false, } - false } fn fill_match_query_columns( diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index dcc5b5c3ac7..a7a127bfcf6 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -3,7 +3,7 @@ use lance_core::{Error, Result}; use serde::{Deserialize, Deserializer, Serialize}; -use std::{env, path::PathBuf}; +use std::{env, path::PathBuf, str::FromStr}; #[cfg(feature = "tokenizer-jieba")] mod jieba; @@ -21,7 +21,7 @@ use lindera::LinderaTokenizerBuilder; use crate::pbold; use crate::pbold::inverted_index_details::DocumentGranularity as PbDocumentGranularity; use crate::scalar::inverted::tokenizer::document_tokenizer::{ - JsonTokenizer, LanceTokenizer, TextTokenizer, + JsonTokenizer, JsonTokenizerMode, LanceTokenizer, TextTokenizer, }; use crate::scalar::inverted::{ InvertedListFormatVersion, default_fts_format_version_for_block_size, @@ -196,6 +196,19 @@ pub struct InvertedIndexParams { /// Index code operators such as `::`, `->`, and `!=`. pub(crate) index_operators: bool, + /// JSON tokenization mode. `None` means the caller did not provide a mode. + /// Existing JSON indexes without this field are interpreted as `SingleDocument`; + /// new JSON indexes default this to `FlattenedSubDocs` during index build. + #[serde(default)] + pub(crate) json_tokenizer_mode: Option, + + /// If true, flattened JSON tokenization avoids cross-array unnesting. + /// This reduces sub-doc explosion for JSON records with multiple sibling + /// arrays by indexing each array independently instead of producing their + /// Cartesian product. Default is false for exact query semantics. + #[serde(default)] + pub(crate) disable_cross_array_unnest: bool, + /// Total memory limit in MiB for the build stage. /// /// This is split evenly across FTS workers at build time. By default Lance @@ -261,6 +274,9 @@ struct RawInvertedIndexParams { split_on_numerics: Option, preserve_original: Option, index_operators: Option, + json_tokenizer_mode: Option, + #[serde(default)] + disable_cross_array_unnest: bool, #[serde(rename = "memory_limit", alias = "worker_memory_limit_mb")] memory_limit_mb: Option, #[serde(rename = "num_workers")] @@ -392,6 +408,10 @@ impl RawInvertedIndexParams { if let Some(index_operators) = self.index_operators { params.index_operators = index_operators; } + if let Some(json_tokenizer_mode) = self.json_tokenizer_mode { + params.json_tokenizer_mode = Some(json_tokenizer_mode); + } + params.disable_cross_array_unnest = self.disable_cross_array_unnest; params.memory_limit_mb = self.memory_limit_mb; params.num_workers = self.num_workers; params.format_version = self.format_version; @@ -427,6 +447,11 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { ), document_granularity: PbDocumentGranularity::from(params.document_granularity) as i32, posting_format_version: Some(params.resolved_format_version().index_version()), + json_tokenizer_mode: params + .json_tokenizer_mode + .filter(|mode| *mode == JsonTokenizerMode::FlattenedSubDocs) + .map(|mode| mode.as_ref().to_string()), + disable_cross_array_unnest: params.disable_cross_array_unnest, }) } } @@ -446,6 +471,12 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { Some(block_size) => validate_block_size(block_size as usize)?, None => LEGACY_BLOCK_SIZE, }, + json_tokenizer_mode: details + .json_tokenizer_mode + .as_deref() + .map(JsonTokenizerMode::from_str) + .transpose()?, + disable_cross_array_unnest: details.disable_cross_array_unnest, ..Self::default() }; params.document_granularity = details.document_granularity.try_into()?; @@ -498,6 +529,12 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { .posting_format_version .map(|version| resolve_fts_format_version(Some(&version.to_string()))) .transpose()?; + params.json_tokenizer_mode = details + .json_tokenizer_mode + .as_deref() + .map(JsonTokenizerMode::from_str) + .transpose()?; + params.disable_cross_array_unnest = details.disable_cross_array_unnest; params.validate()?; Ok(params) } @@ -652,6 +689,8 @@ impl InvertedIndexParams { split_on_numerics: false, preserve_original: false, index_operators: false, + json_tokenizer_mode: None, + disable_cross_array_unnest: false, memory_limit_mb: None, num_workers: None, format_version: None, @@ -749,6 +788,18 @@ impl InvertedIndexParams { self.document_granularity } + /// Set how JSON documents are tokenized by the Lance JSON tokenizer. + pub fn json_tokenizer_mode(mut self, mode: JsonTokenizerMode) -> Self { + self.json_tokenizer_mode = Some(mode); + self + } + + /// Set whether flattened JSON tokenization avoids cross-array unnesting. + pub fn disable_cross_array_unnest(mut self, disable_cross_array_unnest: bool) -> Self { + self.disable_cross_array_unnest = disable_cross_array_unnest; + self + } + /// Set the lexical tokenizer implementation. /// /// Setting this to `"code"` selects the code analyzer defaults. Setting @@ -1039,7 +1090,12 @@ impl InvertedIndexParams { match self.lance_tokenizer { Some(ref t) if t == "text" => Ok(Box::new(TextTokenizer::new(tokenizer))), - Some(ref t) if t == "json" => Ok(Box::new(JsonTokenizer::new(tokenizer))), + Some(ref t) if t == "json" => Ok(Box::new(JsonTokenizer::new( + tokenizer, + self.json_tokenizer_mode + .unwrap_or(JsonTokenizerMode::SingleDocument), + self.disable_cross_array_unnest, + ))), None => Ok(Box::new(TextTokenizer::new(tokenizer))), _ => Err(Error::invalid_input(format!( "unknown lance tokenizer {}", @@ -1587,6 +1643,8 @@ mod tests { code_config: None, document_granularity: PbDocumentGranularity::Row as i32, posting_format_version: None, + json_tokenizer_mode: None, + disable_cross_array_unnest: false, }; let params = InvertedIndexParams::try_from(&old_details).unwrap(); assert_eq!(params.block_size, 128); diff --git a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs index 62dd7b1aa3b..8f40ac02e66 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs @@ -4,8 +4,11 @@ use arrow_schema::{DataType, Field}; use lance_arrow::ARROW_EXT_NAME_KEY; use lance_arrow::json::JSON_EXT_NAME; +use lance_core::Error; use lance_tokenizer::{BoxTokenStream, TextAnalyzer, Token, TokenStream}; +use serde::{Deserialize, Serialize}; use serde_json::Value; +use std::str::FromStr; /// Document type for full text search. #[derive(Debug, Clone)] @@ -14,6 +17,39 @@ pub enum DocType { Json, } +/// Controls how JSON documents are represented inside the inverted index. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum JsonTokenizerMode { + /// Emit one token stream for each source JSON document. + SingleDocument, + /// Flatten arrays into multiple sub-doc token streams for each source JSON document. + FlattenedSubDocs, +} + +impl AsRef for JsonTokenizerMode { + fn as_ref(&self) -> &str { + match self { + Self::SingleDocument => "single_document", + Self::FlattenedSubDocs => "flattened_sub_docs", + } + } +} + +impl FromStr for JsonTokenizerMode { + type Err = Error; + + fn from_str(value: &str) -> std::result::Result { + match value { + "single_document" => Ok(Self::SingleDocument), + "flattened_sub_docs" => Ok(Self::FlattenedSubDocs), + _ => Err(Error::invalid_input(format!( + "unknown JSON tokenizer mode {value:?}; expected 'single_document' or 'flattened_sub_docs'" + ))), + } + } +} + impl AsRef for DocType { fn as_ref(&self) -> &str { match self { @@ -78,10 +114,27 @@ pub trait LanceTokenizer: Send + Sync + std::fmt::Debug { fn token_stream_for_search<'a>(&'a mut self, query_text: &'a str) -> BoxTokenStream<'a>; /// Tokenize document text for index. fn token_stream_for_doc<'a>(&'a mut self, text: &'a str) -> BoxTokenStream<'a>; + /// Tokenize document text into one or more internal inverted-index documents. + fn token_streams_for_doc(&mut self, text: &str) -> lance_core::Result>> { + let mut stream = self.token_stream_for_doc(text); + let mut tokens = Vec::new(); + while let Some(token) = stream.next() { + tokens.push(token.clone()); + } + Ok(vec![tokens]) + } /// Clone the tokenizer. fn box_clone(&self) -> Box; /// Get document type. fn doc_type(&self) -> DocType; + /// Get the JSON tokenization mode, if this tokenizer handles JSON documents. + fn json_tokenizer_mode(&self) -> Option { + None + } + /// Whether flattened JSON tokenization avoids cross-array unnesting. + fn disable_cross_array_unnest(&self) -> bool { + false + } } impl Clone for Box { @@ -128,39 +181,75 @@ impl LanceTokenizer for TextTokenizer { #[derive(Clone)] pub struct JsonTokenizer { tokenizer: TextAnalyzer, + mode: JsonTokenizerMode, + disable_cross_array_unnest: bool, } impl JsonTokenizer { - pub fn new(tokenizer: TextAnalyzer) -> Self { - Self { tokenizer } + pub fn new( + tokenizer: TextAnalyzer, + mode: JsonTokenizerMode, + disable_cross_array_unnest: bool, + ) -> Self { + Self { + tokenizer, + mode, + disable_cross_array_unnest, + } } } impl std::fmt::Debug for JsonTokenizer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "JsonTokenizer") + f.debug_struct("JsonTokenizer") + .field("mode", &self.mode) + .field( + "disable_cross_array_unnest", + &self.disable_cross_array_unnest, + ) + .finish() } } impl LanceTokenizer for JsonTokenizer { fn token_stream_for_search<'a>(&'a mut self, query_text: &'a str) -> BoxTokenStream<'a> { - let tokens = flatten_triplet(query_text, &mut self.tokenizer).unwrap(); + let tokens = flatten_triplet(query_text, self.mode, &mut self.tokenizer).unwrap(); BoxTokenStream::new(TTStream { tokens, index: 0 }) } fn token_stream_for_doc<'a>(&'a mut self, text: &'a str) -> BoxTokenStream<'a> { - let value: Value = match serde_json::from_slice(text.as_bytes()) { - Ok(v) => v, - Err(e) => { - panic!("JSON parse error: {:?}", e); - } - }; - let mut tokens = vec![]; - let mut position = 0; - flatten_json(&value, "", &mut tokens, &mut position, &mut self.tokenizer); + let tokens = self + .token_streams_for_doc(text) + .unwrap() + .into_iter() + .next() + .unwrap_or_default(); BoxTokenStream::new(TTStream { tokens, index: 0 }) } + fn token_streams_for_doc(&mut self, text: &str) -> lance_core::Result>> { + let value: Value = serde_json::from_slice(text.as_bytes()).map_err(|err| { + Error::invalid_input(format!( + "failed to parse JSON document for FTS indexing: {err}" + )) + })?; + + match self.mode { + JsonTokenizerMode::SingleDocument => { + let mut tokens = Vec::new(); + let mut position = 0; + flatten_json(&value, "", &mut tokens, &mut position, &mut self.tokenizer); + Ok(vec![tokens]) + } + JsonTokenizerMode::FlattenedSubDocs => Ok(flatten_json_sub_docs( + &value, + "", + &mut self.tokenizer, + self.disable_cross_array_unnest, + )), + } + } + fn box_clone(&self) -> Box { Box::new(self.clone()) } @@ -168,9 +257,21 @@ impl LanceTokenizer for JsonTokenizer { fn doc_type(&self) -> DocType { DocType::Json } + + fn json_tokenizer_mode(&self) -> Option { + Some(self.mode) + } + + fn disable_cross_array_unnest(&self) -> bool { + self.disable_cross_array_unnest + } } -fn flatten_triplet(text: &str, tokenizer: &mut TextAnalyzer) -> lance_core::Result> { +fn flatten_triplet( + text: &str, + mode: JsonTokenizerMode, + tokenizer: &mut TextAnalyzer, +) -> lance_core::Result> { let mut token_vec = Vec::new(); let mut idx = 0; @@ -184,6 +285,21 @@ fn flatten_triplet(text: &str, tokenizer: &mut TextAnalyzer) -> lance_core::Resu let field = parts[0]; let v_type = parts[1]; let value = parts[2]; + let (field, mut index_tokens) = match mode { + JsonTokenizerMode::SingleDocument => (field.to_string(), Vec::new()), + JsonTokenizerMode::FlattenedSubDocs => normalize_flattened_json_path(field)?, + }; + + for index_token in index_tokens.drain(..) { + token_vec.push(Token { + offset_from: 0, + offset_to: 0, + position: idx, + text: index_token, + position_length: 1, + }); + idx += 1; + } match v_type { "number" | "bool" | "null" => { @@ -220,6 +336,46 @@ fn flatten_triplet(text: &str, tokenizer: &mut TextAnalyzer) -> lance_core::Resu Ok(token_vec) } +fn normalize_flattened_json_path(path: &str) -> lance_core::Result<(String, Vec)> { + let mut normalized = String::with_capacity(path.len()); + let mut index_tokens = Vec::new(); + let mut chars = path.char_indices().peekable(); + + while let Some((_, ch)) = chars.next() { + if ch != '[' { + normalized.push(ch); + continue; + } + + let index_path = normalized.clone(); + let mut array_index = String::new(); + let mut found_right_bracket = false; + for (_, bracket_ch) in chars.by_ref() { + if bracket_ch == ']' { + found_right_bracket = true; + break; + } + array_index.push(bracket_ch); + } + if !found_right_bracket { + return Err(Error::invalid_input(format!( + "missing right bracket in JSON path {path:?}" + ))); + } + if array_index.is_empty() { + return Err(Error::invalid_input(format!( + "empty array index in JSON path {path:?}" + ))); + } + if array_index != "*" { + index_tokens.push(format!("{index_path}$idx,number,{array_index}")); + } + normalized.push('.'); + } + + Ok((normalized, index_tokens)) +} + fn flatten_json( value: &Value, prefix: &str, @@ -277,6 +433,168 @@ fn flatten_json( } } +fn flatten_json_sub_docs( + value: &Value, + prefix: &str, + tokenizer: &mut TextAnalyzer, + disable_cross_array_unnest: bool, +) -> Vec> { + let token_texts = + flatten_json_sub_doc_terms(value, prefix, tokenizer, disable_cross_array_unnest); + token_texts + .into_iter() + .map(|sub_doc| { + sub_doc + .into_iter() + .enumerate() + .map(|(position, text)| Token { + offset_from: 0, + offset_to: 0, + position, + text, + position_length: 1, + }) + .collect() + }) + .collect() +} + +fn flatten_json_sub_doc_terms( + value: &Value, + prefix: &str, + tokenizer: &mut TextAnalyzer, + disable_cross_array_unnest: bool, +) -> Vec> { + match value { + Value::Object(map) => { + let mut non_nested = Vec::new(); + let mut nested: Vec>> = Vec::new(); + + for (key, child) in map { + let child_prefix = if prefix.is_empty() { + key.clone() + } else { + format!("{prefix}.{key}") + }; + let child_terms = flatten_json_sub_doc_terms( + child, + &child_prefix, + tokenizer, + disable_cross_array_unnest, + ); + match child_terms.len() { + 0 => {} + 1 => non_nested.extend(child_terms.into_iter().next().unwrap()), + _ => nested.push(child_terms), + } + } + + match nested.len() { + 0 if non_nested.is_empty() => Vec::new(), + 0 => vec![non_nested], + 1 => nested + .pop() + .unwrap() + .into_iter() + .map(|mut sub_doc| { + sub_doc.extend(non_nested.iter().cloned()); + sub_doc + }) + .collect(), + _ if disable_cross_array_unnest => unnest_json_sub_docs(&nested, &non_nested), + _ => cross_join_json_sub_docs(&nested, &non_nested), + } + } + Value::Array(arr) => { + let mut sub_docs = Vec::new(); + let child_prefix = format!("{prefix}."); + for (array_index, child) in arr.iter().enumerate() { + let mut child_terms = flatten_json_sub_doc_terms( + child, + &child_prefix, + tokenizer, + disable_cross_array_unnest, + ); + for sub_doc in &mut child_terms { + sub_doc.push(format!("{prefix}$idx,number,{array_index}")); + } + sub_docs.extend(child_terms); + } + sub_docs + } + Value::String(text) => { + let mut token_texts = Vec::new(); + let mut tokens = tokenizer.token_stream(text); + while let Some(token) = tokens.next() { + token_texts.push(format!("{prefix},str,{}", token.text)); + } + if token_texts.is_empty() { + Vec::new() + } else { + vec![token_texts] + } + } + _ => { + let value_type = match value { + Value::Null => "null", + Value::Bool(_) => "bool", + Value::Number(_) => "number", + _ => unreachable!(), + }; + vec![vec![format!("{prefix},{value_type},{value}")]] + } + } +} + +fn cross_join_json_sub_docs( + nested: &[Vec>], + non_nested: &[String], +) -> Vec> { + let capacity = nested + .iter() + .map(|sub_docs| sub_docs.len()) + .product::(); + let mut results = Vec::with_capacity(capacity); + let mut current = Vec::new(); + cross_join_json_sub_docs_inner(nested, 0, non_nested, &mut current, &mut results); + results +} + +fn unnest_json_sub_docs(nested: &[Vec>], non_nested: &[String]) -> Vec> { + let capacity = nested.iter().map(|sub_docs| sub_docs.len()).sum::(); + let mut results = Vec::with_capacity(capacity); + for sub_docs in nested { + for child in sub_docs { + let mut sub_doc = child.clone(); + sub_doc.extend(non_nested.iter().cloned()); + results.push(sub_doc); + } + } + results +} + +fn cross_join_json_sub_docs_inner( + nested: &[Vec>], + nested_index: usize, + non_nested: &[String], + current: &mut Vec, + results: &mut Vec>, +) { + if nested_index == nested.len() { + let mut sub_doc = current.clone(); + sub_doc.extend(non_nested.iter().cloned()); + results.push(sub_doc); + return; + } + + for child in &nested[nested_index] { + let old_len = current.len(); + current.extend(child.iter().cloned()); + cross_join_json_sub_docs_inner(nested, nested_index + 1, non_nested, current, results); + current.truncate(old_len); + } +} + struct TTStream { tokens: Vec, index: usize, @@ -304,7 +622,8 @@ impl TokenStream for TTStream { #[cfg(test)] mod tests { use crate::scalar::inverted::tokenizer::document_tokenizer::{ - JsonTokenizer, LanceTokenizer, flatten_json, flatten_triplet, + JsonTokenizer, JsonTokenizerMode, LanceTokenizer, flatten_json, flatten_json_sub_docs, + flatten_triplet, }; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer, Token}; use serde_json::Value; @@ -318,8 +637,11 @@ mod tests { {"c": "e"} ] }"#; - let mut tokenizer = - JsonTokenizer::new(TextAnalyzer::builder(SimpleTokenizer::default()).build()); + let mut tokenizer = JsonTokenizer::new( + TextAnalyzer::builder(SimpleTokenizer::default()).build(), + JsonTokenizerMode::SingleDocument, + false, + ); let mut stream = tokenizer.token_stream_for_doc(text); let mut tokens: Vec = vec![]; @@ -368,7 +690,8 @@ mod tests { fn test_flatten_triplet() { let text = r#"a,number,1;b.c,str,d;b.c,str,e;d,str,hello world;e,number,1.0"#; let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); - let tokens = flatten_triplet(text, &mut tokenizer).unwrap(); + let tokens = + flatten_triplet(text, JsonTokenizerMode::SingleDocument, &mut tokenizer).unwrap(); assert_eq!(tokens.len(), 6); assert_token(&tokens[0], 0, "a,number,1"); @@ -379,6 +702,175 @@ mod tests { assert_token(&tokens[5], 5, "e,number,1.0"); } + #[test] + fn test_flattened_sub_docs_design_example() { + let doc0 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["x","y"]}]}"#); + let doc1 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["y"]},{"bar":"z"}]}"#); + + assert_eq!( + doc0, + vec![ + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,0", + "foo..bar.,str,x", + ]), + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,1", + "foo..bar.,str,y", + ]), + ] + ); + assert_eq!( + doc1, + vec![ + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,0", + "foo..bar.,str,y", + ]), + sorted_tokens(["foo$idx,number,1", "foo..bar,str,z"]), + ] + ); + + let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); + let exact_tokens = flatten_triplet( + "foo[0].bar[0],str,y", + JsonTokenizerMode::FlattenedSubDocs, + &mut tokenizer, + ) + .unwrap(); + assert_token_texts( + &exact_tokens, + &[ + "foo$idx,number,0", + "foo..bar$idx,number,0", + "foo..bar.,str,y", + ], + ); + + let wildcard_tokens = flatten_triplet( + "foo[0].bar[*],str,y", + JsonTokenizerMode::FlattenedSubDocs, + &mut tokenizer, + ) + .unwrap(); + assert_token_texts(&wildcard_tokens, &["foo$idx,number,0", "foo..bar.,str,y"]); + } + + #[test] + fn test_flattened_sub_docs_sibling_array_example() { + let doc0 = flattened_sub_doc_texts( + r#"{"foo":[{"bar":["x","y"]},{"bar":["a","b"]}],"foo2":["u"]}"#, + ); + let doc1 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["y","z"]}],"foo2":["u"]}"#); + + assert_eq!( + doc0, + vec![ + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,0", + "foo..bar.,str,x", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,1", + "foo..bar.,str,y", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + sorted_tokens([ + "foo$idx,number,1", + "foo..bar$idx,number,0", + "foo..bar.,str,a", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + sorted_tokens([ + "foo$idx,number,1", + "foo..bar$idx,number,1", + "foo..bar.,str,b", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + ] + ); + assert_eq!( + doc1, + vec![ + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,0", + "foo..bar.,str,y", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + sorted_tokens([ + "foo$idx,number,0", + "foo..bar$idx,number,1", + "foo..bar.,str,z", + "foo2$idx,number,0", + "foo2.,str,u", + ]), + ] + ); + } + + #[test] + fn test_disable_cross_array_unnest_indexes_arrays_independently() { + let cross_joined = flattened_sub_doc_texts(r#"{"a":["x","y"],"b":["u","v"],"c":1}"#); + let disabled = flattened_sub_doc_texts_with_disable_cross_array_unnest( + r#"{"a":["x","y"],"b":["u","v"],"c":1}"#, + ); + + assert_eq!( + sorted_sub_docs(cross_joined), + sorted_sub_docs(vec![ + sorted_tokens([ + "a$idx,number,0", + "a.,str,x", + "b$idx,number,0", + "b.,str,u", + "c,number,1" + ]), + sorted_tokens([ + "a$idx,number,0", + "a.,str,x", + "b$idx,number,1", + "b.,str,v", + "c,number,1" + ]), + sorted_tokens([ + "a$idx,number,1", + "a.,str,y", + "b$idx,number,0", + "b.,str,u", + "c,number,1" + ]), + sorted_tokens([ + "a$idx,number,1", + "a.,str,y", + "b$idx,number,1", + "b.,str,v", + "c,number,1" + ]), + ]) + ); + assert_eq!( + sorted_sub_docs(disabled), + sorted_sub_docs(vec![ + sorted_tokens(["a$idx,number,0", "a.,str,x", "c,number,1"]), + sorted_tokens(["a$idx,number,1", "a.,str,y", "c,number,1"]), + sorted_tokens(["b$idx,number,0", "b.,str,u", "c,number,1"]), + sorted_tokens(["b$idx,number,1", "b.,str,v", "c,number,1"]), + ]) + ); + } + fn assert_token(token: &Token, position: usize, text: &str) { assert_eq!( token.position, position, @@ -390,4 +882,43 @@ mod tests { "expected text {text} but {token:?}" ); } + + fn flattened_sub_doc_texts(json: &str) -> Vec> { + flattened_sub_doc_texts_with_mode(json, false) + } + + fn flattened_sub_doc_texts_with_disable_cross_array_unnest(json: &str) -> Vec> { + flattened_sub_doc_texts_with_mode(json, true) + } + + fn flattened_sub_doc_texts_with_mode( + json: &str, + disable_cross_array_unnest: bool, + ) -> Vec> { + let value: Value = serde_json::from_str(json).unwrap(); + let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); + flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest) + .into_iter() + .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) + .collect() + } + + fn sorted_tokens(tokens: impl IntoIterator>) -> Vec { + let mut tokens = tokens.into_iter().map(Into::into).collect::>(); + tokens.sort(); + tokens + } + + fn sorted_sub_docs(mut sub_docs: Vec>) -> Vec> { + sub_docs.sort(); + sub_docs + } + + fn assert_token_texts(tokens: &[Token], expected: &[&str]) { + let actual = tokens + .iter() + .map(|token| token.text.as_str()) + .collect::>(); + assert_eq!(actual, expected); + } } diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 2d7c48c72a1..68f5d1c3b1c 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -6457,6 +6457,86 @@ async fn prepare_json_dataset() -> (Dataset, String) { (dataset, json_col) } +#[tokio::test] +async fn test_json_inverted_flattened_sub_doc_array_paths() { + let ids = Arc::new(UInt64Array::from(vec![0, 1])); + let json_col = "json_field".to_string(); + let json_values = Arc::new(StringArray::from(vec![ + r#"{"foo":[{"bar":["x","y"]},{"bar":["a","b"]}],"foo2":["u"]}"#, + r#"{"foo":[{"bar":["y","z"]}],"foo2":["u"]}"#, + ])); + + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata), + ]) + .into(), + vec![ids as ArrayRef, json_values as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(stream, "memory://test/flattened_json_array_paths", None) + .await + .unwrap(); + + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false) + .remove_stop_words(false), + true, + ) + .await + .unwrap(); + + let exact_query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("foo[0].bar[0],str,y".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let exact_batch = dataset + .scan() + .full_text_search(exact_query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let exact_ids = exact_batch["id"].as_primitive::().values(); + assert_eq!(exact_ids, &[1]); + + let wildcard_query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("foo[0].bar[*],str,y".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let wildcard_batch = dataset + .scan() + .full_text_search(wildcard_query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let wildcard_ids = wildcard_batch["id"].as_primitive::().values(); + assert_eq!(wildcard_batch.num_rows(), 2, "ids={wildcard_ids:?}"); + assert!(wildcard_ids.contains(&0), "ids={wildcard_ids:?}"); + assert!(wildcard_ids.contains(&1), "ids={wildcard_ids:?}"); +} + #[tokio::test] async fn test_json_inverted_fuzziness_query() { let (mut dataset, json_col) = prepare_json_dataset().await; @@ -6817,7 +6897,7 @@ async fn test_json_inverted_multimatch_query() { match_queries: vec![ MatchQuery::new("Title,str,harrypotter".to_string()) .with_column(Some(json_col.clone())), - MatchQuery::new("Language,str,english".to_string()) + MatchQuery::new("Language[*],str,english".to_string()) .with_column(Some(json_col.clone())), ], }), @@ -6859,7 +6939,7 @@ async fn test_json_inverted_boolean_query() { should: vec![], must: vec![ FtsQuery::Match( - MatchQuery::new("Language,str,english".to_string()) + MatchQuery::new("Language[*],str,english".to_string()) .with_column(Some(json_col.clone())), ), FtsQuery::Match( diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 3483a476d52..6ae7861a836 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -51,7 +51,9 @@ use crate::{Dataset, index::DatasetIndexInternalExt}; use lance_index::metrics::MetricsCollector; use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::builder::document_input; -use lance_index::scalar::inverted::document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}; +use lance_index::scalar::inverted::document_tokenizer::{ + DocType, JsonTokenizer, JsonTokenizerMode, LanceTokenizer, +}; use lance_index::scalar::inverted::query::{ BoostQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, PhraseQuery, Tokens, collect_query_tokens, has_query_token, uses_fuzzy_expansion, @@ -2052,7 +2054,17 @@ fn tokenizer_for_match_query( let analyzer = TextAnalyzer::from(SimpleTokenizer::default()); match index.tokenizer().doc_type() { DocType::Text => Box::new(TextTokenizer::new(analyzer)), - DocType::Json => Box::new(JsonTokenizer::new(analyzer)), + DocType::Json => { + let index_tokenizer = index.tokenizer(); + let mode = index_tokenizer + .json_tokenizer_mode() + .unwrap_or(JsonTokenizerMode::SingleDocument); + Box::new(JsonTokenizer::new( + analyzer, + mode, + index_tokenizer.disable_cross_array_unnest(), + )) + } } } @@ -2408,6 +2420,28 @@ impl FtsSegmentSelection { } } +fn query_has_concrete_json_array_index(query: &str) -> bool { + query.split(';').any(|triple| { + let path = triple + .split_once(',') + .map(|(path, _)| path) + .unwrap_or(triple); + let mut remaining = path; + while let Some(left_bracket) = remaining.find('[') { + let after_left = &remaining[left_bracket + 1..]; + let Some(right_bracket) = after_left.find(']') else { + return false; + }; + let array_index = &after_left[..right_bracket]; + if !array_index.is_empty() && array_index != "*" { + return true; + } + remaining = &after_left[right_bracket + 1..]; + } + false + }) +} + pub struct FtsIndexMetrics { index_metrics: IndexMetrics, partitions_searched: Count, @@ -2927,13 +2961,21 @@ impl ExecutionPlan for MatchQueryExec { column )))?; let mut tokenizer = tokenizer_for_match_query(first_index, query.fuzziness); + let force_and_operator = tokenizer.json_tokenizer_mode() + == Some(JsonTokenizerMode::FlattenedSubDocs) + && query_has_concrete_json_array_index(&query.terms); + let operator = if force_and_operator { + Operator::And + } else { + query.operator + }; let tokens = collect_query_tokens(&query.terms, &mut tokenizer); record_tokenized_query(&tokenized_query, &tokens); let prepared = if let Some(prepared_query) = preset_prepared_query { Arc::new(PreparedMatch { query: prepared_query, params: Arc::new(params), - operator: query.operator, + operator, }) } else { let base_scorer = match (preset_base_scorer, shared_scorer) { @@ -2954,7 +2996,7 @@ impl ExecutionPlan for MatchQueryExec { &indices, tokens, params, - query.operator, + operator, metrics.as_ref(), base_scorer, ) From b28dc36281748bef109e82e27059012b4f793925 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sun, 21 Jun 2026 23:28:37 +0000 Subject: [PATCH 2/5] test(index): cover JSON cross-object sub-doc matching --- rust/lance/src/dataset/tests/dataset_index.rs | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 68f5d1c3b1c..ad4839d6ddd 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -6537,6 +6537,75 @@ async fn test_json_inverted_flattened_sub_doc_array_paths() { assert!(wildcard_ids.contains(&1), "ids={wildcard_ids:?}"); } +#[tokio::test] +async fn test_json_inverted_flattened_sub_doc_prevents_cross_object_match() { + let ids = Arc::new(UInt64Array::from(vec![0, 1])); + let json_col = "json_field".to_string(); + let json_values = Arc::new(StringArray::from(vec![ + r#"{"cart_id":3234234,"cart":[{"product_type":"sneakers","attributes":{"color":"white"}},{"product_type":"t-shirt","attributes":{"color":"red"}}]}"#, + r#"{"cart_id":3234235,"cart":[{"product_type":"sneakers","attributes":{"color":"red"}}]}"#, + ])); + + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata), + ]) + .into(), + vec![ids as ArrayRef, json_values as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write( + stream, + "memory://test/flattened_json_cross_object_match", + None, + ) + .await + .unwrap(); + + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false) + .remove_stop_words(false), + true, + ) + .await + .unwrap(); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new( + "cart[*].product_type,str,sneakers;cart[*].attributes.color,str,red".to_string(), + ) + .with_column(Some(json_col.clone())) + .with_operator(Operator::And), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_primitive::().values(); + assert_eq!(ids, &[1]); +} + #[tokio::test] async fn test_json_inverted_fuzziness_query() { let (mut dataset, json_col) = prepare_json_dataset().await; From f11104d6e940cfb2ac778258113fe09f36cc6e33 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 4 Sep 2026 03:25:56 +0000 Subject: [PATCH 3/5] refactor(index): simplify flattened JSON sub-doc handling --- .../index/scalar/InvertedIndexParamsTest.java | 8 + python/src/dataset.rs | 13 +- rust/lance-index/src/scalar/inverted.rs | 26 ++ .../src/scalar/inverted/builder.rs | 54 ++- .../src/scalar/inverted/compound.rs | 45 ++- .../src/scalar/inverted/index/flat_search.rs | 90 ++--- .../scalar/inverted/index/inverted_index.rs | 9 +- .../src/scalar/inverted/index/search.rs | 40 +- rust/lance-index/src/scalar/inverted/query.rs | 21 +- .../src/scalar/inverted/tokenizer.rs | 34 +- .../inverted/tokenizer/document_tokenizer.rs | 371 ++++++++---------- rust/lance/src/dataset/mem_wal/index/fts.rs | 202 ++++++---- rust/lance/src/dataset/tests/dataset_index.rs | 172 +++----- rust/lance/src/io/exec/fts.rs | 63 +-- 14 files changed, 538 insertions(+), 610 deletions(-) diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index 59ef98db9e1..2214a4cd843 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -55,6 +55,14 @@ void documentGranularityIsSerialized() { assertEquals("list_element", json.get("document_granularity")); } + @Test + void disableCrossArrayUnnestIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().disableCrossArrayUnnest(true).build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(true, json.get("disable_cross_array_unnest")); + } + @Test void blockSizeIsSerialized() { ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build(); diff --git a/python/src/dataset.rs b/python/src/dataset.rs index ae75f92734d..230b45199b4 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2649,6 +2649,7 @@ impl Dataset { "split_on_numerics", "preserve_original", "index_operators", + "disable_cross_array_unnest", "memory_limit", "num_workers", "format_version", @@ -2682,18 +2683,6 @@ impl Dataset { .block_size(block_size.extract()?) .map_err(|e| PyValueError::new_err(e.to_string()))?; } - if let Some(split_identifiers) = kwargs.get_item("split_identifiers")? { - params = params.split_identifiers(split_identifiers.extract()?); - } - if let Some(split_on_numerics) = kwargs.get_item("split_on_numerics")? { - params = params.split_on_numerics(split_on_numerics.extract()?); - } - if let Some(preserve_original) = kwargs.get_item("preserve_original")? { - params = params.preserve_original(preserve_original.extract()?); - } - if let Some(index_operators) = kwargs.get_item("index_operators")? { - params = params.index_operators(index_operators.extract()?); - } if let Some(disable_cross_array_unnest) = kwargs.get_item("disable_cross_array_unnest")? { diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index b26824e6c27..0ed34b82dca 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -40,6 +40,32 @@ pub use tokenizer::*; use crate::scalar::inverted::query::{FtsSearchParams, Tokens, uses_fuzzy_expansion}; +pub(crate) fn collapse_scored_rows( + rows: impl IntoIterator, + limit: usize, +) -> Vec<(u64, f32)> { + let mut scores_by_row_id = HashMap::new(); + for (row_id, score) in rows { + scores_by_row_id + .entry(row_id) + .and_modify(|existing| { + if score > *existing { + *existing = score; + } + }) + .or_insert(score); + } + + let mut rows = scores_by_row_id.into_iter().collect::>(); + rows.sort_unstable_by(|(left_id, left_score), (right_id, right_score)| { + right_score + .total_cmp(left_score) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(rows.len().min(limit)); + rows +} + /// Canonical token vocabulary and BM25 statistics for one indexed query leaf. /// /// Keeping these values together prevents a search path from expanding one diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 7681843b7b9..b6e47b6be8c 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -32,6 +32,7 @@ use lance_select::RowSetOps; use object_store::path::Path; use roaring::RoaringBitmap; use smallvec::SmallVec; +use std::borrow::Cow; use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; @@ -183,19 +184,13 @@ impl InvertedIndexBuilder { /// Constructed as `(fragment_id as u64) << 32`. /// When provided, ensures that generated IDs belong to the specified fragment. pub fn from_existing_index( - mut params: InvertedIndexParams, + params: InvertedIndexParams, store: Option>, partitions: Vec, token_set_format: TokenSetFormat, fragment_mask: Option, deleted_fragments: RoaringBitmap, ) -> Self { - if (store.is_some() || !partitions.is_empty()) - && params.lance_tokenizer.as_deref() == Some("json") - && params.json_tokenizer_mode.is_none() - { - params.json_tokenizer_mode = Some(JsonTokenizerMode::SingleDocument); - } let format_version = params.resolved_format_version(); Self { params, @@ -211,15 +206,10 @@ impl InvertedIndexBuilder { } } - fn configure_json_tokenizer_mode_for_new_data(&mut self, doc_type: DocType) { + fn infer_lance_tokenizer(&mut self, doc_type: DocType) { if self.params.lance_tokenizer.is_none() { self.params.lance_tokenizer = Some(doc_type.as_ref().to_string()); } - if self.params.lance_tokenizer.as_deref() == Some("json") - && self.params.json_tokenizer_mode.is_none() - { - self.params.json_tokenizer_mode = Some(JsonTokenizerMode::FlattenedSubDocs); - } } pub fn with_posting_tail_codec(mut self, posting_tail_codec: PostingTailCodec) -> Self { @@ -265,7 +255,12 @@ impl InvertedIndexBuilder { // infer lance_tokenizer based on document type let field = schema.column_with_name(doc_col).expect_ok()?.1; let doc_type = DocType::try_from(field)?; - self.configure_json_tokenizer_mode_for_new_data(doc_type); + self.infer_lance_tokenizer(doc_type); + if self.params.lance_tokenizer.as_deref() == Some("json") + && self.params.json_tokenizer_mode.is_none() + { + self.params.json_tokenizer_mode = Some(JsonTokenizerMode::FlattenedSubDocs); + } let new_data = document_input(new_data, doc_col)?; @@ -296,7 +291,7 @@ impl InvertedIndexBuilder { let field = schema.column_with_name(doc_col).expect_ok()?.1; let doc_type = DocType::try_from(field)?; - self.configure_json_tokenizer_mode_for_new_data(doc_type); + self.infer_lance_tokenizer(doc_type); let mut files = self .merge_existing_segments(dest_store, old_segments, old_data_filter.as_ref()) @@ -1593,12 +1588,14 @@ impl IndexWorker { doc_index: &[u32], ) -> Result<()> { let doc = match document { - DocumentSource::Text(doc) => doc.to_string(), - DocumentSource::StringList(elements) => Self::materialize_string_list(elements), + DocumentSource::Text(doc) => Cow::Borrowed(doc), + DocumentSource::StringList(elements) => { + Cow::Owned(Self::materialize_string_list(elements)) + } }; self.total_doc_length += doc.len(); let with_position = self.has_position(); - let sub_docs = self.tokenizer.token_streams_for_doc(&doc)?; + let sub_docs = self.tokenizer.token_streams_for_doc(doc.as_ref())?; for tokens in sub_docs { self.process_tokenized_doc(row_id, tokens, with_position, doc_index) .await?; @@ -1610,7 +1607,7 @@ impl IndexWorker { async fn process_tokenized_doc( &mut self, row_id: u64, - mut tokens: Vec, + tokens: Vec, with_position: bool, doc_index: &[u32], ) -> Result<()> { @@ -1620,19 +1617,20 @@ impl IndexWorker { let doc_id = self.builder.docs.len() as u32; let mut token_num: u32 = 0; let mut posting_memory_delta = 0i64; + if self.token_ids.capacity() < self.last_token_count { + self.token_ids + .reserve(self.last_token_count - self.token_ids.capacity()); + } + self.token_ids.clear(); + if with_position { - if self.token_ids.capacity() < self.last_token_count { - self.token_ids - .reserve(self.last_token_count - self.token_ids.capacity()); - } - self.token_ids.clear(); let builder = &mut self.builder; let token_ids = &mut self.token_ids; let memory_size = &mut self.memory_size; let posting_tail_codec = builder.posting_tail_codec; let block_size = builder.block_size; - for token in &mut tokens { + for token in tokens { let position = Self::checked_token_position(row_id, token.position)?; let token_id = builder.tokens.get_or_add(&token.text); if token_id as usize == builder.posting_lists.len() { @@ -1666,12 +1664,6 @@ impl IndexWorker { token_num += 1; } } else { - if self.token_ids.capacity() < self.last_token_count { - self.token_ids - .reserve(self.last_token_count - self.token_ids.capacity()); - } - self.token_ids.clear(); - for token in tokens { let token_id = self.builder.tokens.get_or_add(&token.text); self.token_ids.push(token_id); diff --git a/rust/lance-index/src/scalar/inverted/compound.rs b/rust/lance-index/src/scalar/inverted/compound.rs index b0c4b485218..f09608ca8b2 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -16,7 +16,7 @@ use lance_select::RowAddrMask; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; use super::{ - InvertedIndex, PreparedBm25Query, + InvertedIndex, PreparedBm25Query, collapse_scored_rows, document_tokenizer::{DocType, JsonTokenizer, JsonTokenizerMode, LanceTokenizer}, documents::{ CachedRowAddressOrder, DocId, DocLengths, DocVisibility, OrderedRowAddressProjection, @@ -26,6 +26,7 @@ use super::{ prepare_bm25_query, query::{ FtsQuery, FtsSearchParams, MatchQuery, Operator, PhraseQuery, Tokens, collect_query_tokens, + effective_json_query_operator, }, scorer::MemBM25Scorer, tokenizer::document_tokenizer::TextTokenizer, @@ -3794,13 +3795,13 @@ pub(super) fn tokenize_leaf( let index_tokenizer = index.tokenizer(); match index_tokenizer.doc_type() { DocType::Text => Box::new(TextTokenizer::new(analyzer)) as Box, - DocType::Json => Box::new(JsonTokenizer::new( - analyzer, - index_tokenizer - .json_tokenizer_mode() - .unwrap_or(JsonTokenizerMode::SingleDocument), - index_tokenizer.disable_cross_array_unnest(), - )) as Box, + DocType::Json => Box::new( + JsonTokenizer::new(analyzer).with_mode( + index_tokenizer + .json_tokenizer_mode() + .unwrap_or(JsonTokenizerMode::SingleDocument), + ), + ) as Box, } } else { index.tokenizer() @@ -3839,6 +3840,11 @@ async fn prepare_compound_query( for leaf in leaf_queries { let effective_params = leaf.effective_params(params); let tokens = tokenize_leaf(first_index, &leaf, &effective_params); + let operator = effective_json_query_operator( + first_index.params().json_tokenizer_mode, + &tokens, + leaf.operator(), + ); let prepared = match &prepared_match { Some(prepared) => prepared.clone(), None => Arc::new( @@ -3855,7 +3861,7 @@ async fn prepare_compound_query( leaves.push(PreparedLeaf { query: prepared, params: Arc::new(effective_params), - operator: leaf.operator(), + operator, }); } Ok((plan, leaves)) @@ -4360,6 +4366,9 @@ async fn compound_search_impl( if limit == 0 { return Ok((Vec::new(), Vec::new())); } + let should_collapse_rows = indices.first().is_some_and(|index| { + index.params().json_tokenizer_mode == Some(JsonTokenizerMode::FlattenedSubDocs) + }); let (plan, leaves) = prepare_compound_query( indices, query, @@ -4375,7 +4384,12 @@ async fn compound_search_impl( if let Some(score_floor) = initial_score_floor { competitive_score.raise(checked_score(score_floor, "initial compound score floor")?); } - let mut collector = TopKCollector::with_competitive_score(limit, competitive_score); + let collector_limit = if should_collapse_rows { + usize::MAX + } else { + limit + }; + let mut collector = TopKCollector::with_competitive_score(collector_limit, competitive_score); for (segment_ordinal, index) in indices.iter().enumerate() { let loads = @@ -4436,8 +4450,15 @@ async fn compound_search_impl( } } - let rows = collector.into_rows(); - Ok(rows.into_iter().map(|row| (row.row_id, row.score)).unzip()) + let rows = collector + .into_rows() + .into_iter() + .map(|row| (row.row_id, row.score)); + if should_collapse_rows { + Ok(collapse_scored_rows(rows, limit).into_iter().unzip()) + } else { + Ok(rows.unzip()) + } } #[cfg(test)] diff --git a/rust/lance-index/src/scalar/inverted/index/flat_search.rs b/rust/lance-index/src/scalar/inverted/index/flat_search.rs index 97ba79a0ee4..d4f1f242f75 100644 --- a/rust/lance-index/src/scalar/inverted/index/flat_search.rs +++ b/rust/lance-index/src/scalar/inverted/index/flat_search.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use super::*; +use crate::scalar::inverted::collapse_scored_rows; pub fn doc_index_storage_column(rank: usize) -> String { format!("{DOC_INDEX_STORAGE_PREFIX}{rank}") @@ -137,28 +138,28 @@ pub(super) fn document_matches_flat_query( return Ok(has_query_token(document, tokenizer, query_tokens)); }; - let mut document_positions = (0..query_tokens.len()) - .map(|_| Vec::new()) - .collect::>(); - let mut stream = tokenizer.token_stream_for_doc(document); - while let Some(token) = stream.next() { - let position = u32::try_from(token.position).map_err(|_| { - Error::invalid_input(format!( - "flat FTS token position exceeds u32: {}", - token.position - )) - })?; - for (query_index, positions) in document_positions.iter_mut().enumerate() { - if query_tokens.get_token(query_index) == token.text { - positions.push(position); + for tokens in tokenizer.token_streams_for_doc(document)? { + let mut document_positions = (0..query_tokens.len()) + .map(|_| Vec::new()) + .collect::>(); + for token in tokens { + let position = u32::try_from(token.position).map_err(|_| { + Error::invalid_input(format!( + "flat FTS token position exceeds u32: {}", + token.position + )) + })?; + for (query_index, positions) in document_positions.iter_mut().enumerate() { + if query_tokens.get_token(query_index) == token.text { + positions.push(position); + } } } + if phrase_matches_positions(query_tokens, &document_positions, slop) { + return Ok(true); + } } - Ok(phrase_matches_positions( - query_tokens, - &document_positions, - slop, - )) + Ok(false) } pub(super) const FLAT_ALL_TOKENS_COL: &str = "all_tokens"; @@ -572,40 +573,16 @@ pub(super) fn initialize_scorer( MemBM25Scorer::new(total_tokens, num_docs, token_counts_map) } -fn deduplicate_scored_rows( - row_ids: Vec, - scores: Vec, - limit: usize, -) -> (Vec, Vec) { - let mut scores_by_row_id: HashMap = HashMap::with_capacity(row_ids.len()); - for (row_id, score) in row_ids.into_iter().zip(scores) { - scores_by_row_id - .entry(row_id) - .and_modify(|existing| { - if score > *existing { - *existing = score; - } - }) - .or_insert(score); - } - - let mut scored_rows = scores_by_row_id.into_iter().collect::>(); - scored_rows.sort_unstable_by(|(left_row_id, left_score), (right_row_id, right_score)| { - right_score - .total_cmp(left_score) - .then_with(|| left_row_id.cmp(right_row_id)) - }); - scored_rows.truncate(scored_rows.len().min(limit)); - scored_rows.into_iter().unzip() -} - -fn deduplicate_fts_batch(batch: RecordBatch, limit: usize) -> Result { +fn collapse_flattened_rows(batch: RecordBatch) -> Result { let row_ids = batch[ROW_ID].as_primitive::().values().to_vec(); let scores = batch[SCORE_COL] .as_primitive::() .values() .to_vec(); - let (row_ids, scores) = deduplicate_scored_rows(row_ids, scores, limit); + let (row_ids, scores): (Vec<_>, Vec<_>) = + collapse_scored_rows(row_ids.into_iter().zip(scores), usize::MAX) + .into_iter() + .unzip(); Ok(RecordBatch::try_new( FTS_SCHEMA.clone(), vec![ @@ -623,7 +600,6 @@ pub(super) fn flat_bm25_score( operator: Operator, boost: f32, phrase_slop: Option, - require_all_query_tokens: bool, ) -> Result { let mut row_ids_builder = UInt64Builder::with_capacity(counted_input.num_rows()); let mut scores_builder = Float32Builder::with_capacity(counted_input.num_rows()); @@ -709,16 +685,12 @@ pub(super) fn flat_bm25_score( } let doc_norm = K1 * (1.0 - B + B * num_tokens_in_doc as f32 / scorer.avg_doc_length()); let mut score = 0.0; - let mut has_all_query_tokens = true; for (token, freq) in query_tokens.into_iter().zip(query_token_counts) { let freq = freq as f32; - if freq == 0.0 { - has_all_query_tokens = false; - } let idf = idf(scorer.num_docs_containing_token(token), scorer.num_docs()); score += idf * (freq * (K1 + 1.0) / (freq + doc_norm)); } - if score > 0.0 && (!require_all_query_tokens || has_all_query_tokens) { + if score > 0.0 { row_ids_builder.append_value(row_id); if let Some(builder) = doc_indices_builder.as_mut() { for input_doc_index in &input_doc_indices { @@ -911,11 +883,8 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( let query_tokens = Arc::new(collect_query_tokens(&query, &mut tokenizer)); let should_deduplicate_rows = tokenizer.json_tokenizer_mode() == Some(JsonTokenizerMode::FlattenedSubDocs); - let require_all_query_tokens = should_deduplicate_rows - && query_tokens - .as_ref() - .into_iter() - .any(|token| token.contains("$idx,number,")); + let operator = + effective_json_query_operator(tokenizer.json_tokenizer_mode(), &query_tokens, operator); // A query that tokenizes to no terms (e.g. only stop words) has no // searchable content and matches nothing. Return early rather than @@ -985,10 +954,9 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( operator, boost, phrase_slop, - require_all_query_tokens, )?; if should_deduplicate_rows { - scores = deduplicate_fts_batch(scores, usize::MAX)?; + scores = collapse_flattened_rows(scores)?; } // Finally we emit batches according to the target batch size diff --git a/rust/lance-index/src/scalar/inverted/index/inverted_index.rs b/rust/lance-index/src/scalar/inverted/index/inverted_index.rs index dd4151c64b2..98c72ee6d67 100644 --- a/rust/lance-index/src/scalar/inverted/index/inverted_index.rs +++ b/rust/lance-index/src/scalar/inverted/index/inverted_index.rs @@ -85,10 +85,15 @@ impl InvertedIndex { } fn to_builder_with_offset(&self, fragment_mask: Option) -> InvertedIndexBuilder { + let mut params = self.params.clone(); + if params.lance_tokenizer.as_deref() == Some("json") && params.json_tokenizer_mode.is_none() + { + params.json_tokenizer_mode = Some(JsonTokenizerMode::SingleDocument); + } if self.is_legacy() { // for legacy format, we re-create the index in the new format InvertedIndexBuilder::from_existing_index( - self.params.clone(), + params, None, Vec::new(), self.token_set_format, @@ -111,7 +116,7 @@ impl InvertedIndex { }; InvertedIndexBuilder::from_existing_index( - self.params.clone(), + params, Some(self.store.clone()), partitions, self.token_set_format, diff --git a/rust/lance-index/src/scalar/inverted/index/search.rs b/rust/lance-index/src/scalar/inverted/index/search.rs index 6149d2b5d7c..52e10e2d4d1 100644 --- a/rust/lance-index/src/scalar/inverted/index/search.rs +++ b/rust/lance-index/src/scalar/inverted/index/search.rs @@ -3,6 +3,7 @@ use super::partition::FuzzyAutomaton; use super::*; +use crate::scalar::inverted::collapse_scored_rows; const LANCE_FTS_REUSE_PREPARED_SCORER_ENV: &str = "LANCE_FTS_REUSE_PREPARED_SCORER"; @@ -639,7 +640,15 @@ impl InvertedIndex { .await }?; if should_deduplicate_rows { - Ok(deduplicate_scored_documents(documents, requested_limit)) + Ok(collapse_scored_rows( + documents + .into_iter() + .map(|document| (document.row_id, document.score.0)), + requested_limit, + ) + .into_iter() + .map(|(row_id, score)| ScoredDoc::new(row_id, score)) + .collect()) } else { Ok(documents) } @@ -1186,35 +1195,6 @@ impl InvertedIndex { Ok(resolved_documents) } } - -fn deduplicate_scored_documents(documents: Vec, limit: usize) -> Vec { - let mut scores_by_row_id: HashMap = HashMap::with_capacity(documents.len()); - for document in documents { - let score = document.score.0; - scores_by_row_id - .entry(document.row_id) - .and_modify(|existing| { - if score > *existing { - *existing = score; - } - }) - .or_insert(score); - } - - let mut scored_documents = scores_by_row_id - .into_iter() - .map(|(row_id, score)| ScoredDoc::new(row_id, score)) - .collect::>(); - scored_documents.sort_unstable_by(|left, right| { - right - .score - .cmp(&left.score) - .then_with(|| left.row_id.cmp(&right.row_id)) - }); - scored_documents.truncate(scored_documents.len().min(limit)); - scored_documents -} - #[cfg(test)] mod tests { use super::*; diff --git a/rust/lance-index/src/scalar/inverted/query.rs b/rust/lance-index/src/scalar/inverted/query.rs index e0a19056856..f7c56c6a909 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use crate::scalar::inverted::DocumentGranularity; -use crate::scalar::inverted::document_tokenizer::DocType; +use crate::scalar::inverted::document_tokenizer::{DocType, JsonTokenizerMode}; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; use lance_core::{Error, Result}; use serde::ser::SerializeMap; @@ -899,6 +899,25 @@ pub fn collect_query_tokens(text: &str, tokenizer: &mut Box) Tokens::with_positions(tokens, positions, token_type) } +#[doc(hidden)] +pub fn effective_json_query_operator( + mode: Option, + query_tokens: &Tokens, + operator: Operator, +) -> Operator { + if mode == Some(JsonTokenizerMode::FlattenedSubDocs) + && query_tokens.into_iter().any(|token| { + token + .split_once(',') + .is_some_and(|(path, value)| path.ends_with("$idx") && value.starts_with("number,")) + }) + { + Operator::And + } else { + operator + } +} + pub fn has_query_token( text: &str, tokenizer: &mut Box, diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index a7a127bfcf6..c0b96cdcd91 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -788,13 +788,19 @@ impl InvertedIndexParams { self.document_granularity } - /// Set how JSON documents are tokenized by the Lance JSON tokenizer. - pub fn json_tokenizer_mode(mut self, mode: JsonTokenizerMode) -> Self { - self.json_tokenizer_mode = Some(mode); - self - } - - /// Set whether flattened JSON tokenization avoids cross-array unnesting. + /// Set whether flattened JSON tokenization indexes sibling arrays independently. + /// + /// This avoids the Cartesian product of sibling arrays at the cost of accuracy + /// for queries that constrain values across those arrays. + /// + /// # Examples + /// + /// ``` + /// use lance_index::scalar::InvertedIndexParams; + /// + /// let params = InvertedIndexParams::default().disable_cross_array_unnest(true); + /// assert!(params.build().is_ok()); + /// ``` pub fn disable_cross_array_unnest(mut self, disable_cross_array_unnest: bool) -> Self { self.disable_cross_array_unnest = disable_cross_array_unnest; self @@ -1090,12 +1096,14 @@ impl InvertedIndexParams { match self.lance_tokenizer { Some(ref t) if t == "text" => Ok(Box::new(TextTokenizer::new(tokenizer))), - Some(ref t) if t == "json" => Ok(Box::new(JsonTokenizer::new( - tokenizer, - self.json_tokenizer_mode - .unwrap_or(JsonTokenizerMode::SingleDocument), - self.disable_cross_array_unnest, - ))), + Some(ref t) if t == "json" => Ok(Box::new( + JsonTokenizer::new(tokenizer) + .with_mode( + self.json_tokenizer_mode + .unwrap_or(JsonTokenizerMode::SingleDocument), + ) + .with_disable_cross_array_unnest(self.disable_cross_array_unnest), + )), None => Ok(Box::new(TextTokenizer::new(tokenizer))), _ => Err(Error::invalid_input(format!( "unknown lance tokenizer {}", diff --git a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs index 8f40ac02e66..4672e0f4b62 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs @@ -131,10 +131,6 @@ pub trait LanceTokenizer: Send + Sync + std::fmt::Debug { fn json_tokenizer_mode(&self) -> Option { None } - /// Whether flattened JSON tokenization avoids cross-array unnesting. - fn disable_cross_array_unnest(&self) -> bool { - false - } } impl Clone for Box { @@ -186,17 +182,27 @@ pub struct JsonTokenizer { } impl JsonTokenizer { - pub fn new( - tokenizer: TextAnalyzer, - mode: JsonTokenizerMode, - disable_cross_array_unnest: bool, - ) -> Self { + pub fn new(tokenizer: TextAnalyzer) -> Self { Self { tokenizer, - mode, - disable_cross_array_unnest, + mode: JsonTokenizerMode::SingleDocument, + disable_cross_array_unnest: false, } } + + #[doc(hidden)] + pub fn with_mode(mut self, mode: JsonTokenizerMode) -> Self { + self.mode = mode; + self + } + + pub(crate) fn with_disable_cross_array_unnest( + mut self, + disable_cross_array_unnest: bool, + ) -> Self { + self.disable_cross_array_unnest = disable_cross_array_unnest; + self + } } impl std::fmt::Debug for JsonTokenizer { @@ -261,10 +267,6 @@ impl LanceTokenizer for JsonTokenizer { fn json_tokenizer_mode(&self) -> Option { Some(self.mode) } - - fn disable_cross_array_unnest(&self) -> bool { - self.disable_cross_array_unnest - } } fn flatten_triplet( @@ -439,9 +441,8 @@ fn flatten_json_sub_docs( tokenizer: &mut TextAnalyzer, disable_cross_array_unnest: bool, ) -> Vec> { - let token_texts = - flatten_json_sub_doc_terms(value, prefix, tokenizer, disable_cross_array_unnest); - token_texts + flatten_json_sub_doc_terms(value, prefix, tokenizer, disable_cross_array_unnest) + .sub_docs .into_iter() .map(|sub_doc| { sub_doc @@ -459,12 +460,17 @@ fn flatten_json_sub_docs( .collect() } +struct FlattenedJsonSubDocs { + sub_docs: Vec>, + has_array: bool, +} + fn flatten_json_sub_doc_terms( value: &Value, prefix: &str, tokenizer: &mut TextAnalyzer, disable_cross_array_unnest: bool, -) -> Vec> { +) -> FlattenedJsonSubDocs { match value { Value::Object(map) => { let mut non_nested = Vec::new(); @@ -482,27 +488,46 @@ fn flatten_json_sub_doc_terms( tokenizer, disable_cross_array_unnest, ); - match child_terms.len() { + match child_terms.sub_docs.len() { 0 => {} - 1 => non_nested.extend(child_terms.into_iter().next().unwrap()), - _ => nested.push(child_terms), + 1 if !child_terms.has_array => { + non_nested.extend(child_terms.sub_docs.into_iter().flatten()) + } + _ => nested.push(child_terms.sub_docs), } } - match nested.len() { - 0 if non_nested.is_empty() => Vec::new(), - 0 => vec![non_nested], - 1 => nested - .pop() - .unwrap() - .into_iter() - .map(|mut sub_doc| { - sub_doc.extend(non_nested.iter().cloned()); - sub_doc - }) - .collect(), - _ if disable_cross_array_unnest => unnest_json_sub_docs(&nested, &non_nested), - _ => cross_join_json_sub_docs(&nested, &non_nested), + if nested.is_empty() { + return FlattenedJsonSubDocs { + sub_docs: if non_nested.is_empty() { + Vec::new() + } else { + vec![non_nested] + }, + has_array: false, + }; + } + if nested.len() == 1 { + return FlattenedJsonSubDocs { + sub_docs: nested + .into_iter() + .flatten() + .map(|mut sub_doc| { + sub_doc.extend(non_nested.iter().cloned()); + sub_doc + }) + .collect(), + has_array: true, + }; + } + let sub_docs = if disable_cross_array_unnest { + unnest_json_sub_docs(&nested, &non_nested) + } else { + cross_join_json_sub_docs(&nested, &non_nested) + }; + FlattenedJsonSubDocs { + sub_docs, + has_array: true, } } Value::Array(arr) => { @@ -515,12 +540,15 @@ fn flatten_json_sub_doc_terms( tokenizer, disable_cross_array_unnest, ); - for sub_doc in &mut child_terms { + for sub_doc in &mut child_terms.sub_docs { sub_doc.push(format!("{prefix}$idx,number,{array_index}")); } - sub_docs.extend(child_terms); + sub_docs.extend(child_terms.sub_docs); + } + FlattenedJsonSubDocs { + sub_docs, + has_array: true, } - sub_docs } Value::String(text) => { let mut token_texts = Vec::new(); @@ -528,21 +556,27 @@ fn flatten_json_sub_doc_terms( while let Some(token) = tokens.next() { token_texts.push(format!("{prefix},str,{}", token.text)); } - if token_texts.is_empty() { - Vec::new() - } else { - vec![token_texts] + FlattenedJsonSubDocs { + sub_docs: if token_texts.is_empty() { + Vec::new() + } else { + vec![token_texts] + }, + has_array: false, } } - _ => { - let value_type = match value { - Value::Null => "null", - Value::Bool(_) => "bool", - Value::Number(_) => "number", - _ => unreachable!(), - }; - vec![vec![format!("{prefix},{value_type},{value}")]] - } + Value::Null => FlattenedJsonSubDocs { + sub_docs: vec![vec![format!("{prefix},null,null")]], + has_array: false, + }, + Value::Bool(value) => FlattenedJsonSubDocs { + sub_docs: vec![vec![format!("{prefix},bool,{value}")]], + has_array: false, + }, + Value::Number(value) => FlattenedJsonSubDocs { + sub_docs: vec![vec![format!("{prefix},number,{value}")]], + has_array: false, + }, } } @@ -637,11 +671,8 @@ mod tests { {"c": "e"} ] }"#; - let mut tokenizer = JsonTokenizer::new( - TextAnalyzer::builder(SimpleTokenizer::default()).build(), - JsonTokenizerMode::SingleDocument, - false, - ); + let mut tokenizer = + JsonTokenizer::new(TextAnalyzer::builder(SimpleTokenizer::default()).build()); let mut stream = tokenizer.token_stream_for_doc(text); let mut tokens: Vec = vec![]; @@ -703,35 +734,40 @@ mod tests { } #[test] - fn test_flattened_sub_docs_design_example() { - let doc0 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["x","y"]}]}"#); - let doc1 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["y"]},{"bar":"z"}]}"#); - - assert_eq!( - doc0, - vec![ - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,0", - "foo..bar.,str,x", - ]), - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,1", - "foo..bar.,str,y", - ]), - ] + fn test_flattened_sub_docs_examples() { + assert_sub_docs( + r#"{"foo":[{"bar":["x","y"]}]}"#, + false, + &[ + "foo$idx,number,0;foo..bar$idx,number,0;foo..bar.,str,x", + "foo$idx,number,0;foo..bar$idx,number,1;foo..bar.,str,y", + ], ); - assert_eq!( - doc1, - vec![ - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,0", - "foo..bar.,str,y", - ]), - sorted_tokens(["foo$idx,number,1", "foo..bar,str,z"]), - ] + assert_sub_docs( + r#"{"foo":[{"bar":["y"]},{"bar":"z"}]}"#, + false, + &[ + "foo$idx,number,0;foo..bar$idx,number,0;foo..bar.,str,y", + "foo$idx,number,1;foo..bar,str,z", + ], + ); + assert_sub_docs( + r#"{"foo":[{"bar":["x","y"]},{"bar":["a","b"]}],"foo2":["u"]}"#, + false, + &[ + "foo$idx,number,0;foo..bar$idx,number,0;foo..bar.,str,x;foo2$idx,number,0;foo2.,str,u", + "foo$idx,number,0;foo..bar$idx,number,1;foo..bar.,str,y;foo2$idx,number,0;foo2.,str,u", + "foo$idx,number,1;foo..bar$idx,number,0;foo..bar.,str,a;foo2$idx,number,0;foo2.,str,u", + "foo$idx,number,1;foo..bar$idx,number,1;foo..bar.,str,b;foo2$idx,number,0;foo2.,str,u", + ], + ); + assert_sub_docs( + r#"{"foo":[{"bar":["y","z"]}],"foo2":["u"]}"#, + false, + &[ + "foo$idx,number,0;foo..bar$idx,number,0;foo..bar.,str,y;foo2$idx,number,0;foo2.,str,u", + "foo$idx,number,0;foo..bar$idx,number,1;foo..bar.,str,z;foo2$idx,number,0;foo2.,str,u", + ], ); let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); @@ -760,114 +796,35 @@ mod tests { } #[test] - fn test_flattened_sub_docs_sibling_array_example() { - let doc0 = flattened_sub_doc_texts( - r#"{"foo":[{"bar":["x","y"]},{"bar":["a","b"]}],"foo2":["u"]}"#, - ); - let doc1 = flattened_sub_doc_texts(r#"{"foo":[{"bar":["y","z"]}],"foo2":["u"]}"#); - - assert_eq!( - doc0, - vec![ - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,0", - "foo..bar.,str,x", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,1", - "foo..bar.,str,y", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - sorted_tokens([ - "foo$idx,number,1", - "foo..bar$idx,number,0", - "foo..bar.,str,a", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - sorted_tokens([ - "foo$idx,number,1", - "foo..bar$idx,number,1", - "foo..bar.,str,b", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - ] - ); - assert_eq!( - doc1, - vec![ - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,0", - "foo..bar.,str,y", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - sorted_tokens([ - "foo$idx,number,0", - "foo..bar$idx,number,1", - "foo..bar.,str,z", - "foo2$idx,number,0", - "foo2.,str,u", - ]), - ] - ); - } - - #[test] - fn test_disable_cross_array_unnest_indexes_arrays_independently() { - let cross_joined = flattened_sub_doc_texts(r#"{"a":["x","y"],"b":["u","v"],"c":1}"#); - let disabled = flattened_sub_doc_texts_with_disable_cross_array_unnest( + fn test_disable_cross_array_unnest() { + assert_sub_docs( r#"{"a":["x","y"],"b":["u","v"],"c":1}"#, + false, + &[ + "a$idx,number,0;a.,str,x;b$idx,number,0;b.,str,u;c,number,1", + "a$idx,number,0;a.,str,x;b$idx,number,1;b.,str,v;c,number,1", + "a$idx,number,1;a.,str,y;b$idx,number,0;b.,str,u;c,number,1", + "a$idx,number,1;a.,str,y;b$idx,number,1;b.,str,v;c,number,1", + ], ); - - assert_eq!( - sorted_sub_docs(cross_joined), - sorted_sub_docs(vec![ - sorted_tokens([ - "a$idx,number,0", - "a.,str,x", - "b$idx,number,0", - "b.,str,u", - "c,number,1" - ]), - sorted_tokens([ - "a$idx,number,0", - "a.,str,x", - "b$idx,number,1", - "b.,str,v", - "c,number,1" - ]), - sorted_tokens([ - "a$idx,number,1", - "a.,str,y", - "b$idx,number,0", - "b.,str,u", - "c,number,1" - ]), - sorted_tokens([ - "a$idx,number,1", - "a.,str,y", - "b$idx,number,1", - "b.,str,v", - "c,number,1" - ]), - ]) + assert_sub_docs( + r#"{"a":["x","y"],"b":["u","v"],"c":1}"#, + true, + &[ + "a$idx,number,0;a.,str,x;c,number,1", + "a$idx,number,1;a.,str,y;c,number,1", + "b$idx,number,0;b.,str,u;c,number,1", + "b$idx,number,1;b.,str,v;c,number,1", + ], ); - assert_eq!( - sorted_sub_docs(disabled), - sorted_sub_docs(vec![ - sorted_tokens(["a$idx,number,0", "a.,str,x", "c,number,1"]), - sorted_tokens(["a$idx,number,1", "a.,str,y", "c,number,1"]), - sorted_tokens(["b$idx,number,0", "b.,str,u", "c,number,1"]), - sorted_tokens(["b$idx,number,1", "b.,str,v", "c,number,1"]), - ]) + assert_sub_docs( + r#"{"a":["x"],"b":["u","v"],"c":null}"#, + true, + &[ + "a$idx,number,0;a.,str,x;c,null,null", + "b$idx,number,0;b.,str,u;c,null,null", + "b$idx,number,1;b.,str,v;c,null,null", + ], ); } @@ -883,24 +840,21 @@ mod tests { ); } - fn flattened_sub_doc_texts(json: &str) -> Vec> { - flattened_sub_doc_texts_with_mode(json, false) - } - - fn flattened_sub_doc_texts_with_disable_cross_array_unnest(json: &str) -> Vec> { - flattened_sub_doc_texts_with_mode(json, true) - } - - fn flattened_sub_doc_texts_with_mode( - json: &str, - disable_cross_array_unnest: bool, - ) -> Vec> { + fn assert_sub_docs(json: &str, disable_cross_array_unnest: bool, expected: &[&str]) { let value: Value = serde_json::from_str(json).unwrap(); let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); - flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest) - .into_iter() - .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) - .collect() + let mut actual = + flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest) + .into_iter() + .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) + .collect::>(); + actual.sort(); + let mut expected = expected + .iter() + .map(|sub_doc| sorted_tokens(sub_doc.split(';'))) + .collect::>(); + expected.sort(); + assert_eq!(actual, expected); } fn sorted_tokens(tokens: impl IntoIterator>) -> Vec { @@ -909,11 +863,6 @@ mod tests { tokens } - fn sorted_sub_docs(mut sub_docs: Vec>) -> Vec> { - sub_docs.sort(); - sub_docs - } - fn assert_token_texts(tokens: &[Token], expected: &[&str]) { let actual = tokens .iter() diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 50ab8cc3b36..1fe5b128df2 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -58,10 +58,14 @@ use lance_bitpacking::{BitPacker, BitPacker4x}; use lance_core::datatypes::Schema as LanceSchema; use lance_core::{Error, Result}; use lance_index::scalar::InvertedIndexParams; -use lance_index::scalar::inverted::query::{FtsQuery, Operator, Tokens}; -use lance_index::scalar::inverted::tokenizer::document_tokenizer::{DocType, LanceTokenizer}; +use lance_index::scalar::inverted::query::{ + FtsQuery, Operator, Tokens, effective_json_query_operator, +}; +use lance_index::scalar::inverted::tokenizer::document_tokenizer::{ + DocType, JsonTokenizerMode, LanceTokenizer, +}; use lance_index::scalar::inverted::{DocSet, MemBM25Scorer, Scorer, TokenSet}; -use lance_tokenizer::TokenStream; +use lance_tokenizer::{Token, TokenStream}; use rayon::prelude::*; use rustc_hash::{FxHashMap, FxHashSet}; @@ -995,6 +999,7 @@ pub struct FtsMemIndex { resolved_field: OnceLock, tokenizer_pool: Arc, + flattened_sub_docs: bool, /// Writer-only tokenizer slot. Held under a Mutex purely so `insert` /// can take `&self`. Single-writer assumption means this is uncontested. writer_tokenizer: Mutex>, @@ -1130,6 +1135,7 @@ impl QueryLocalFtsIndex { params: self.inner.params.clone(), resolved_field, tokenizer_pool: self.inner.tokenizer_pool.clone(), + flattened_sub_docs: self.inner.flattened_sub_docs, writer_tokenizer: Mutex::new(self.inner.tokenizer_pool.acquire()), state: ArcSwap::from(IndexState::empty()), freeze_threshold_rows: self.inner.freeze_threshold_rows, @@ -1274,6 +1280,8 @@ impl FtsMemIndex { pool: TokenizerPool, background_maintenance: bool, ) -> Self { + let flattened_sub_docs = + pool.template.json_tokenizer_mode() == Some(JsonTokenizerMode::FlattenedSubDocs); let writer_tokenizer = pool.template.box_clone(); Self { field_id, @@ -1281,6 +1289,7 @@ impl FtsMemIndex { params, resolved_field: OnceLock::new(), tokenizer_pool: Arc::new(pool), + flattened_sub_docs, writer_tokenizer: Mutex::new(writer_tokenizer), state: ArcSwap::from(IndexState::empty()), freeze_threshold_rows: Self::DEFAULT_FREEZE_THRESHOLD_ROWS, @@ -1516,37 +1525,39 @@ impl FtsMemIndex { let preserve_zero_token_documents = self.params.get_document_granularity().is_list_element(); let mut index_document = |key: DocumentKey, text: &str| -> Result<()> { - let document_position = document_position_start + documents.len() as u64; - let (num_tokens, retained_term) = match allowed_terms { - Some(allowed_terms) => index_text_filtered( - text, - document_position, - tokenizer, - &mut term_builders, - allowed_terms, - )?, - None => ( - index_text(text, document_position, tokenizer, &mut term_builders)?, - false, - ), - }; - let belongs_in_corpus = preserve_zero_token_documents || num_tokens > 0; - if allowed_terms.is_some() && belongs_in_corpus { - query_local_corpus_doc_count = query_local_corpus_doc_count - .checked_add(1) - .ok_or_else(|| Error::internal("query-local FTS document count overflow"))?; - query_local_corpus_total_tokens = query_local_corpus_total_tokens - .checked_add(num_tokens as u64) - .ok_or_else(|| Error::internal("query-local FTS total token count overflow"))?; - } - let retain_document = if allowed_terms.is_some() { - retained_term - } else { - belongs_in_corpus - }; - if retain_document { - documents.push(DocumentMetadata { key, num_tokens }); - total_tokens += num_tokens as u64; + for (sub_doc_index, tokens) in tokenizer + .token_streams_for_doc(text)? + .into_iter() + .enumerate() + { + let document_position = document_position_start + documents.len() as u64; + let (num_tokens, retained_term) = + index_tokens(tokens, document_position, &mut term_builders, allowed_terms)?; + let belongs_in_corpus = preserve_zero_token_documents || num_tokens > 0; + if allowed_terms.is_some() && belongs_in_corpus { + query_local_corpus_doc_count = + query_local_corpus_doc_count.checked_add(1).ok_or_else(|| { + Error::internal("query-local FTS document count overflow") + })?; + query_local_corpus_total_tokens = query_local_corpus_total_tokens + .checked_add(num_tokens as u64) + .ok_or_else(|| { + Error::internal("query-local FTS total token count overflow") + })?; + } + let retain_document = if allowed_terms.is_some() { + retained_term + } else { + belongs_in_corpus + }; + if retain_document { + let mut key = key.clone(); + if self.flattened_sub_docs { + key.doc_index.push(sub_doc_index as u32); + } + documents.push(DocumentMetadata { key, num_tokens }); + total_tokens += num_tokens as u64; + } } Ok(()) }; @@ -1667,8 +1678,11 @@ impl FtsMemIndex { } let st = index.state.load_full(); let tokens = index.analyze_for_search(&query.terms); + let operator = index.effective_match_operator(&tokens, query.operator); let rows = index - .search_match_with_scorer(&st, &tokens, query.operator, scorer) + .collapse_sub_docs( + index.search_match_with_scorer(&st, &tokens, operator, scorer), + ) .into_iter() .map(|entry| (entry.row_position, entry.score)) .collect(); @@ -1678,7 +1692,9 @@ impl FtsMemIndex { let st = index.state.load_full(); let tokens = index.analyze_for_search(&query.terms); let rows = index - .search_phrase_with_scorer(&st, &tokens, query.slop, scorer) + .collapse_sub_docs( + index.search_phrase_with_scorer(&st, &tokens, query.slop, scorer), + ) .into_iter() .map(|entry| (entry.row_position, entry.score)) .collect(); @@ -1809,7 +1825,8 @@ impl FtsMemIndex { pub fn search(&self, term: &str) -> Vec { let st = self.state.load_full(); let tokens = self.analyze_for_search(term); - self.search_match(&st, &tokens, Operator::Or, None, true, true) + let operator = self.effective_match_operator(&tokens, Operator::Or); + self.collapse_sub_docs(self.search_match(&st, &tokens, operator, None, true, true)) } /// Search for documents containing an exact phrase, optionally allowing @@ -1817,7 +1834,7 @@ impl FtsMemIndex { pub fn search_phrase(&self, phrase: &str, slop: u32) -> Vec { let st = self.state.load_full(); let tokens = self.analyze_for_search(phrase); - self.search_phrase_tokens(&st, &tokens, slop, true) + self.collapse_sub_docs(self.search_phrase_tokens(&st, &tokens, slop, true)) } /// Freeze the current mutable tail into an immutable partition, so a @@ -1855,7 +1872,14 @@ impl FtsMemIndex { ) -> Vec { let st = self.state.load_full(); let tokens = self.tokenize_for_search(query); - self.search_fuzzy_tokens(&st, &tokens, fuzziness, 0, max_expansions, true) + self.collapse_sub_docs(self.search_fuzzy_tokens( + &st, + &tokens, + fuzziness, + 0, + max_expansions, + true, + )) } /// BM25 OR-search over the query tokens, scored with one corpus-wide @@ -2319,14 +2343,27 @@ impl FtsMemIndex { boost, } => { let tokens = self.analyze_for_search(query); - let mut results = - self.search_match(st, &tokens, *operator, limit, include_tail, tail_skip); + let operator = self.effective_match_operator(&tokens, *operator); + let search_limit = if self.flattened_sub_docs { None } else { limit }; + let mut results = self.collapse_sub_docs(self.search_match( + st, + &tokens, + operator, + search_limit, + include_tail, + tail_skip, + )); apply_boost(&mut results, *boost); results } FtsQueryExpr::Phrase { query, slop, boost } => { let tokens = self.analyze_for_search(query); - let mut results = self.search_phrase_tokens(st, &tokens, *slop, include_tail); + let mut results = self.collapse_sub_docs(self.search_phrase_tokens( + st, + &tokens, + *slop, + include_tail, + )); apply_boost(&mut results, *boost); results } @@ -2338,14 +2375,14 @@ impl FtsMemIndex { boost, } => { let tokens = self.tokenize_for_search(query); - let mut results = self.search_fuzzy_tokens( + let mut results = self.collapse_sub_docs(self.search_fuzzy_tokens( st, &tokens, *fuzziness, *prefix_length, *max_expansions, include_tail, - ); + )); apply_boost(&mut results, *boost); results } @@ -2512,6 +2549,39 @@ impl FtsMemIndex { Tokens::with_positions(tokens, positions, DocType::Text) } + fn effective_match_operator(&self, tokens: &Tokens, operator: Operator) -> Operator { + effective_json_query_operator( + self.flattened_sub_docs + .then_some(JsonTokenizerMode::FlattenedSubDocs), + tokens, + operator, + ) + } + + fn collapse_sub_docs(&self, entries: Vec) -> Vec { + if !self.flattened_sub_docs { + return entries; + } + + let mut rows = HashMap::::with_capacity(entries.len()); + for mut entry in entries { + if let Some(doc_index) = entry.doc_index.as_mut() { + doc_index.pop(); + if doc_index.is_empty() { + entry.doc_index = None; + } + } + rows.entry(entry.key()) + .and_modify(|existing| { + if entry.score > existing.score { + existing.score = entry.score; + } + }) + .or_insert(entry); + } + rows.into_values().collect() + } + // ------------------------------------------------------------------ // Flush to Lance inverted index format // ------------------------------------------------------------------ @@ -2578,10 +2648,15 @@ impl FtsMemIndex { let mut original_to_doc_id: HashMap = HashMap::with_capacity(entries.len()); for (key, num_tokens) in &entries { - let doc_id = if !key.doc_index.is_empty() { - docs.append_with_doc_index(key.row_position, *num_tokens, &key.doc_index)? + let doc_index = if self.flattened_sub_docs { + &key.doc_index[..key.doc_index.len() - 1] } else { + key.doc_index.as_slice() + }; + let doc_id = if doc_index.is_empty() { docs.append(key.row_position, *num_tokens) + } else { + docs.append_with_doc_index(key.row_position, *num_tokens, doc_index)? }; original_to_doc_id.insert(key.clone(), doc_id); } @@ -2761,40 +2836,15 @@ impl BatchTermBuilder { } } -fn index_text( - text: &str, - document_position: u64, - tokenizer: &mut dyn LanceTokenizer, - term_builders: &mut FxHashMap, BatchTermBuilder>, -) -> Result { - index_text_with_predicate(text, document_position, tokenizer, term_builders, |_| true) - .map(|(num_tokens, _)| num_tokens) -} - -fn index_text_filtered( - text: &str, +fn index_tokens( + tokens: Vec, document_position: u64, - tokenizer: &mut dyn LanceTokenizer, term_builders: &mut FxHashMap, BatchTermBuilder>, - allowed_terms: &FxHashSet, + allowed_terms: Option<&FxHashSet>, ) -> Result<(u32, bool)> { - index_text_with_predicate(text, document_position, tokenizer, term_builders, |term| { - allowed_terms.contains(term) - }) -} - -#[inline] -fn index_text_with_predicate( - text: &str, - document_position: u64, - tokenizer: &mut dyn LanceTokenizer, - term_builders: &mut FxHashMap, BatchTermBuilder>, - mut retain_term: impl FnMut(&str) -> bool, -) -> Result<(u32, bool)> { - let mut stream = tokenizer.token_stream_for_doc(text); let mut num_tokens = 0u32; let mut retained_term = false; - while let Some(token) = stream.next() { + for token in tokens { let position = u32::try_from(token.position).map_err(|_| { Error::invalid_input(format!( "token position overflow for document_position={document_position}: token_position={}", @@ -2802,7 +2852,7 @@ fn index_text_with_predicate( )) })?; let term = token.text.as_str(); - if retain_term(term) { + if allowed_terms.is_none_or(|allowed_terms| allowed_terms.contains(term)) { retained_term = true; if let Some(builder) = term_builders.get_mut(term) { builder.observe(document_position, position); diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index ad4839d6ddd..5a6f5517027 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -6458,117 +6458,34 @@ async fn prepare_json_dataset() -> (Dataset, String) { } #[tokio::test] -async fn test_json_inverted_flattened_sub_doc_array_paths() { - let ids = Arc::new(UInt64Array::from(vec![0, 1])); +async fn test_json_inverted_flattened_sub_docs() { let json_col = "json_field".to_string(); - let json_values = Arc::new(StringArray::from(vec![ - r#"{"foo":[{"bar":["x","y"]},{"bar":["a","b"]}],"foo2":["u"]}"#, - r#"{"foo":[{"bar":["y","z"]}],"foo2":["u"]}"#, - ])); - let mut metadata = HashMap::new(); metadata.insert( ARROW_EXT_NAME_KEY.to_string(), ARROW_JSON_EXT_NAME.to_string(), ); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("id", DataType::UInt64, false), - Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata), - ]) - .into(), - vec![ids as ArrayRef, json_values as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(stream, "memory://test/flattened_json_array_paths", None) - .await - .unwrap(); - - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .stem(false) - .remove_stop_words(false), - true, + let schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata), + ])); + let make_batch = |id, json| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(vec![id])) as ArrayRef, + Arc::new(StringArray::from(vec![json])) as ArrayRef, + ], ) - .await - .unwrap(); - - let exact_query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("foo[0].bar[0],str,y".to_string()).with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let exact_batch = dataset - .scan() - .full_text_search(exact_query) .unwrap() - .try_into_batch() - .await - .unwrap(); - let exact_ids = exact_batch["id"].as_primitive::().values(); - assert_eq!(exact_ids, &[1]); - - let wildcard_query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("foo[0].bar[*],str,y".to_string()).with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, }; - let wildcard_batch = dataset - .scan() - .full_text_search(wildcard_query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let wildcard_ids = wildcard_batch["id"].as_primitive::().values(); - assert_eq!(wildcard_batch.num_rows(), 2, "ids={wildcard_ids:?}"); - assert!(wildcard_ids.contains(&0), "ids={wildcard_ids:?}"); - assert!(wildcard_ids.contains(&1), "ids={wildcard_ids:?}"); -} -#[tokio::test] -async fn test_json_inverted_flattened_sub_doc_prevents_cross_object_match() { - let ids = Arc::new(UInt64Array::from(vec![0, 1])); - let json_col = "json_field".to_string(); - let json_values = Arc::new(StringArray::from(vec![ - r#"{"cart_id":3234234,"cart":[{"product_type":"sneakers","attributes":{"color":"white"}},{"product_type":"t-shirt","attributes":{"color":"red"}}]}"#, - r#"{"cart_id":3234235,"cart":[{"product_type":"sneakers","attributes":{"color":"red"}}]}"#, - ])); - - let mut metadata = HashMap::new(); - metadata.insert( - ARROW_EXT_NAME_KEY.to_string(), - ARROW_JSON_EXT_NAME.to_string(), + let indexed_batch = make_batch( + 0, + r#"{"cart":[{"product_type":"sneakers","attributes":{"color":"white"}},{"product_type":"t-shirt","attributes":{"color":"blue"}},{"product_type":"hat","attributes":{"color":"red"}}]}"#, ); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("id", DataType::UInt64, false), - Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata), - ]) - .into(), - vec![ids as ArrayRef, json_values as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write( - stream, - "memory://test/flattened_json_cross_object_match", - None, - ) - .await - .unwrap(); + let stream = RecordBatchIterator::new(vec![Ok(indexed_batch)], schema.clone()); + let mut dataset = Dataset::write(stream, "memory://", None).await.unwrap(); dataset .create_index( @@ -6584,26 +6501,43 @@ async fn test_json_inverted_flattened_sub_doc_prevents_cross_object_match() { .await .unwrap(); - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new( - "cart[*].product_type,str,sneakers;cart[*].attributes.color,str,red".to_string(), - ) - .with_column(Some(json_col.clone())) - .with_operator(Operator::And), + let appended_batch = make_batch( + 1, + r#"{"cart":[{"product_type":"t-shirt","attributes":{"color":"blue"}},{"product_type":"sneakers","attributes":{"color":"red"}}]}"#, + ); + let stream = RecordBatchIterator::new(vec![Ok(appended_batch)], schema); + dataset.append(stream, None).await.unwrap(); + + let cases = [ + ("cart[1].attributes.color,str,red", Operator::Or, vec![1]), + ("cart[*].attributes.color,str,red", Operator::Or, vec![0, 1]), + ( + "cart[*].product_type,str,sneakers;cart[*].attributes.color,str,red", + Operator::And, + vec![1], ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let ids = batch["id"].as_primitive::().values(); - assert_eq!(ids, &[1]); + ]; + for (terms, operator, expected_ids) in cases { + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new(terms.to_string()) + .with_column(Some(json_col.clone())) + .with_operator(operator), + ), + limit: Some(2), + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut ids = batch["id"].as_primitive::().values().to_vec(); + ids.sort_unstable(); + assert_eq!(ids, expected_ids, "query={terms}"); + } } #[tokio::test] diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 6ae7861a836..b36e28f375b 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -56,7 +56,7 @@ use lance_index::scalar::inverted::document_tokenizer::{ }; use lance_index::scalar::inverted::query::{ BoostQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, PhraseQuery, Tokens, - collect_query_tokens, has_query_token, uses_fuzzy_expansion, + collect_query_tokens, effective_json_query_operator, uses_fuzzy_expansion, }; use lance_index::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_index::scalar::inverted::{ @@ -2059,11 +2059,7 @@ fn tokenizer_for_match_query( let mode = index_tokenizer .json_tokenizer_mode() .unwrap_or(JsonTokenizerMode::SingleDocument); - Box::new(JsonTokenizer::new( - analyzer, - mode, - index_tokenizer.disable_cross_array_unnest(), - )) + Box::new(JsonTokenizer::new(analyzer).with_mode(mode)) } } } @@ -2420,28 +2416,6 @@ impl FtsSegmentSelection { } } -fn query_has_concrete_json_array_index(query: &str) -> bool { - query.split(';').any(|triple| { - let path = triple - .split_once(',') - .map(|(path, _)| path) - .unwrap_or(triple); - let mut remaining = path; - while let Some(left_bracket) = remaining.find('[') { - let after_left = &remaining[left_bracket + 1..]; - let Some(right_bracket) = after_left.find(']') else { - return false; - }; - let array_index = &after_left[..right_bracket]; - if !array_index.is_empty() && array_index != "*" { - return true; - } - remaining = &after_left[right_bracket + 1..]; - } - false - }) -} - pub struct FtsIndexMetrics { index_metrics: IndexMetrics, partitions_searched: Count, @@ -2961,15 +2935,12 @@ impl ExecutionPlan for MatchQueryExec { column )))?; let mut tokenizer = tokenizer_for_match_query(first_index, query.fuzziness); - let force_and_operator = tokenizer.json_tokenizer_mode() - == Some(JsonTokenizerMode::FlattenedSubDocs) - && query_has_concrete_json_array_index(&query.terms); - let operator = if force_and_operator { - Operator::And - } else { - query.operator - }; let tokens = collect_query_tokens(&query.terms, &mut tokenizer); + let operator = effective_json_query_operator( + tokenizer.json_tokenizer_mode(), + &tokens, + query.operator, + ); record_tokenized_query(&tokenized_query, &tokens); let prepared = if let Some(prepared_query) = preset_prepared_query { Arc::new(PreparedMatch { @@ -3075,8 +3046,13 @@ fn document_matches_query( query_tokens: &Tokens, operator: Operator, ) -> bool { - match operator { - Operator::Or => has_query_token(text, tokenizer, query_tokens), + let Ok(sub_docs) = tokenizer.token_streams_for_doc(text) else { + return false; + }; + sub_docs.into_iter().any(|tokens| match operator { + Operator::Or => tokens + .iter() + .any(|token| query_tokens.contains(&token.text)), Operator::And => { let mut remaining_positions = (0..query_tokens.len()) .map(|index| query_tokens.position(index)) @@ -3084,8 +3060,7 @@ fn document_matches_query( if remaining_positions.is_empty() { return false; } - let mut stream = tokenizer.token_stream_for_doc(text); - while let Some(token) = stream.next() { + for token in tokens { for index in 0..query_tokens.len() { if token.text == query_tokens.get_token(index) { remaining_positions.remove(&query_tokens.position(index)); @@ -3097,7 +3072,7 @@ fn document_matches_query( } false } - } + }) } impl DisplayAs for FlatMatchFilterExec { @@ -3331,6 +3306,11 @@ impl FlatMatchFilterExec { } }; let query_tokens = Arc::new(collect_query_tokens(&query.terms, &mut tokenizer)); + let query_operator = effective_json_query_operator( + tokenizer.json_tokenizer_mode(), + &query_tokens, + query.operator, + ); record_tokenized_query(&tokenized_query, &query_tokens); let baseline = BaselineMetrics::new(&metrics_set, partition); @@ -3341,7 +3321,6 @@ impl FlatMatchFilterExec { let mut tokenizer = tokenizer.box_clone(); let elapsed_compute = elapsed_compute.clone(); let resolved_field = resolved_field.clone(); - let query_operator = query.operator; async move { let batch = batch_result?; let _t = elapsed_compute.timer(); From 243b5e1020d9f1b51bcf22fc3c8d701fa6ae3a5b Mon Sep 17 00:00:00 2001 From: xuanyili Date: Thu, 10 Sep 2026 02:19:19 +0000 Subject: [PATCH 4/5] feat(index): limit flattened JSON sub-docs per row --- .../index/scalar/InvertedIndexParams.java | 40 ++++ .../scalar/MaxSubDocsPerRowExceedAction.java | 32 +++ .../index/scalar/InvertedIndexParamsTest.java | 15 ++ protos/index_old.proto | 13 ++ python/python/lance/dataset.py | 8 + python/python/tests/test_scalar_index.py | 4 + python/src/dataset.rs | 14 ++ rust/lance-index/src/scalar.rs | 2 +- .../src/scalar/inverted/tokenizer.rs | 147 +++++++++++- .../inverted/tokenizer/document_tokenizer.rs | 210 +++++++++++++----- 10 files changed, 428 insertions(+), 57 deletions(-) create mode 100644 java/src/main/java/org/lance/index/scalar/MaxSubDocsPerRowExceedAction.java diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index 852ba18ffcd..bf2ec0ad534 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -58,6 +58,8 @@ public static final class Builder { private Boolean prefixOnly; private Integer blockSize = 128; private Boolean disableCrossArrayUnnest; + private Long maxSubDocsPerRow; + private MaxSubDocsPerRowExceedAction maxSubDocsPerRowExceedAction; private Boolean splitIdentifiers; private Boolean splitOnNumerics; private Boolean preserveOriginal; @@ -321,6 +323,37 @@ public Builder disableCrossArrayUnnest(boolean disableCrossArrayUnnest) { return this; } + /** + * Limit the number of flattened sub-documents emitted for one JSON row. + * + *

If unset, the number of sub-documents is unlimited. + * + * @param maxSubDocsPerRow maximum sub-documents per row, must be positive + * @return this builder + * @throws IllegalArgumentException if {@code maxSubDocsPerRow} is not positive + */ + public Builder maxSubDocsPerRow(long maxSubDocsPerRow) { + if (maxSubDocsPerRow <= 0) { + throw new IllegalArgumentException("maxSubDocsPerRow must be positive"); + } + this.maxSubDocsPerRow = maxSubDocsPerRow; + return this; + } + + /** + * Configure the action taken when {@link #maxSubDocsPerRow(long)} is exceeded. + * + *

The default is {@link MaxSubDocsPerRowExceedAction#FAIL}. + * + * @param action action to take when the limit is exceeded + * @return this builder + */ + public Builder maxSubDocsPerRowExceedAction(MaxSubDocsPerRowExceedAction action) { + this.maxSubDocsPerRowExceedAction = + Objects.requireNonNull(action, "maxSubDocsPerRowExceedAction must not be null"); + return this; + } + /** * Configure whether code identifiers are split into subwords. * @@ -521,6 +554,13 @@ public ScalarIndexParams build() { if (disableCrossArrayUnnest != null) { params.put("disable_cross_array_unnest", disableCrossArrayUnnest); } + if (maxSubDocsPerRow != null) { + params.put("max_sub_docs_per_row", maxSubDocsPerRow); + } + if (maxSubDocsPerRowExceedAction != null) { + params.put( + "max_sub_docs_per_row_exceed_action", maxSubDocsPerRowExceedAction.toRustString()); + } if (splitIdentifiers != null) { params.put("split_identifiers", splitIdentifiers); } diff --git a/java/src/main/java/org/lance/index/scalar/MaxSubDocsPerRowExceedAction.java b/java/src/main/java/org/lance/index/scalar/MaxSubDocsPerRowExceedAction.java new file mode 100644 index 00000000000..aa6200439ca --- /dev/null +++ b/java/src/main/java/org/lance/index/scalar/MaxSubDocsPerRowExceedAction.java @@ -0,0 +1,32 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.index.scalar; + +/** Action taken when a JSON row exceeds the configured flattened sub-document limit. */ +public enum MaxSubDocsPerRowExceedAction { + /** Abort index ingestion with an error. */ + FAIL("fail"), + /** Omit the source row from the index and continue ingestion. */ + SKIP_ROW("skip_row"); + + private final String rustValue; + + MaxSubDocsPerRowExceedAction(String rustValue) { + this.rustValue = rustValue; + } + + String toRustString() { + return rustValue; + } +} diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index 2214a4cd843..81a68d2fc0c 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -63,6 +63,21 @@ void disableCrossArrayUnnestIsSerialized() { assertEquals(true, json.get("disable_cross_array_unnest")); } + @Test + void maxSubDocsPerRowOptionsAreSerialized() { + for (MaxSubDocsPerRowExceedAction action : MaxSubDocsPerRowExceedAction.values()) { + ScalarIndexParams params = + InvertedIndexParams.builder() + .maxSubDocsPerRow(128) + .maxSubDocsPerRowExceedAction(action) + .build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(128, ((Number) json.get("max_sub_docs_per_row")).longValue()); + assertEquals(action.toRustString(), json.get("max_sub_docs_per_row_exceed_action")); + } + } + @Test void blockSizeIsSerialized() { ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build(); diff --git a/protos/index_old.proto b/protos/index_old.proto index 0f7a84fe3cd..bcb5eeed9f7 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -57,6 +57,11 @@ message InvertedIndexDetails { LIST_ELEMENT = 1; } + enum MaxSubDocsPerRowExceedAction { + FAIL = 0; + SKIP_ROW = 1; + } + message CodeTokenizerConfig { /* Split one lexical identifier into subwords, e.g. getUserName -> * get/user/name. @@ -121,4 +126,12 @@ message InvertedIndexDetails { // If true, avoid cross-array unnesting during flattened JSON tokenization. // The default false value preserves exact Cartesian-product semantics. bool disable_cross_array_unnest = 17; + /* Maximum flattened sub-documents emitted for one JSON row. Absence means + * unlimited. + */ + optional uint64 max_sub_docs_per_row = 18; + /* Action taken when max_sub_docs_per_row is exceeded. FAIL aborts index + * ingestion; SKIP_ROW omits the source row from the index. + */ + MaxSubDocsPerRowExceedAction max_sub_docs_per_row_exceed_action = 19; } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 9f4dc3cb093..f65825aa1aa 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3726,6 +3726,14 @@ def create_scalar_index( producing their Cartesian product. This reduces index build memory for records with multiple arrays but can sacrifice result accuracy for queries that constrain values across those arrays. + max_sub_docs_per_row: int, optional + This is for the ``INVERTED`` index on JSON columns. Maximum number of + flattened sub-documents one source row may produce. If unset, the + number is unlimited. + max_sub_docs_per_row_exceed_action: str, default "fail" + Action when ``max_sub_docs_per_row`` is exceeded. ``"fail"`` aborts + index ingestion with tuning guidance. ``"skip_row"`` omits the source + row from the index and continues. base_tokenizer: str, default "simple" This is for the ``INVERTED`` index. The base tokenizer to use. The value can be: diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 6d7bfba1747..6db248f19bc 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -6033,9 +6033,13 @@ def test_json_inverted_match_query(tmp_path): lower_case=True, remove_stop_words=True, disable_cross_array_unnest=True, + max_sub_docs_per_row=128, + max_sub_docs_per_row_exceed_action="skip_row", ) details = dataset.describe_indices()[0].details assert details["disable_cross_array_unnest"] is True + assert details["max_sub_docs_per_row"] == 128 + assert details["max_sub_docs_per_row_exceed_action"] == "skip_row" # Test match query with token exceeding max_token_length results = dataset.to_table( diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 230b45199b4..1373e71b93f 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2650,6 +2650,8 @@ impl Dataset { "preserve_original", "index_operators", "disable_cross_array_unnest", + "max_sub_docs_per_row", + "max_sub_docs_per_row_exceed_action", "memory_limit", "num_workers", "format_version", @@ -2689,6 +2691,18 @@ impl Dataset { params = params .disable_cross_array_unnest(disable_cross_array_unnest.extract()?); } + if let Some(max_sub_docs_per_row) = kwargs.get_item("max_sub_docs_per_row")? { + params = params + .max_sub_docs_per_row(max_sub_docs_per_row.extract()?) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + } + if let Some(action) = kwargs.get_item("max_sub_docs_per_row_exceed_action")? { + let action: String = action.extract()?; + params = + params.max_sub_docs_per_row_exceed_action(action.parse().map_err( + |err: lance_core::Error| PyValueError::new_err(err.to_string()), + )?); + } if let Some(memory_limit) = kwargs.get_item("memory_limit")? { params = params.memory_limit_mb(memory_limit.extract()?); } diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 6dcfc9018e0..2088cd3cd25 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -47,7 +47,7 @@ pub mod seed; pub mod zoned; pub mod zonemap; -pub use inverted::tokenizer::InvertedIndexParams; +pub use inverted::tokenizer::{InvertedIndexParams, MaxSubDocsPerRowExceedAction}; /// Convert a `Vec<`[`lance_index_core::scalar::IndexFile`]`>` to a /// `Vec<`[`lance_table::format::IndexFile`]`>`. diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index c0b96cdcd91..9f14bdf53b0 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -20,6 +20,7 @@ use lindera::LinderaTokenizerBuilder; use crate::pbold; use crate::pbold::inverted_index_details::DocumentGranularity as PbDocumentGranularity; +use crate::pbold::inverted_index_details::MaxSubDocsPerRowExceedAction as PbMaxSubDocsPerRowExceedAction; use crate::scalar::inverted::tokenizer::document_tokenizer::{ JsonTokenizer, JsonTokenizerMode, LanceTokenizer, TextTokenizer, }; @@ -100,6 +101,54 @@ impl From for PbDocumentGranularity { } } +/// Action taken when a JSON row exceeds [`InvertedIndexParams::max_sub_docs_per_row`]. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, Default)] +#[serde(rename_all = "snake_case")] +pub enum MaxSubDocsPerRowExceedAction { + /// Abort index ingestion with an error. + #[default] + Fail, + /// Omit the source row from the index and continue ingestion. + SkipRow, +} + +impl FromStr for MaxSubDocsPerRowExceedAction { + type Err = Error; + + fn from_str(value: &str) -> Result { + match value { + "fail" => Ok(Self::Fail), + "skip_row" => Ok(Self::SkipRow), + _ => Err(Error::invalid_input(format!( + "unknown max_sub_docs_per_row_exceed_action {value:?}; expected 'fail' or 'skip_row'" + ))), + } + } +} + +impl TryFrom for MaxSubDocsPerRowExceedAction { + type Error = Error; + + fn try_from(value: i32) -> Result { + match PbMaxSubDocsPerRowExceedAction::try_from(value) { + Ok(PbMaxSubDocsPerRowExceedAction::Fail) => Ok(Self::Fail), + Ok(PbMaxSubDocsPerRowExceedAction::SkipRow) => Ok(Self::SkipRow), + Err(_) => Err(Error::invalid_input(format!( + "unknown max_sub_docs_per_row_exceed_action value {value}" + ))), + } + } +} + +impl From for PbMaxSubDocsPerRowExceedAction { + fn from(value: MaxSubDocsPerRowExceedAction) -> Self { + match value { + MaxSubDocsPerRowExceedAction::Fail => Self::Fail, + MaxSubDocsPerRowExceedAction::SkipRow => Self::SkipRow, + } + } +} + /// Tokenizer configs #[derive(Debug, Clone, Serialize, PartialEq)] pub struct InvertedIndexParams { @@ -209,6 +258,16 @@ pub struct InvertedIndexParams { #[serde(default)] pub(crate) disable_cross_array_unnest: bool, + /// Maximum flattened sub-documents one JSON row may produce. + /// + /// `None` leaves the count unlimited. + #[serde(default)] + pub(crate) max_sub_docs_per_row: Option, + + /// Action taken when a JSON row exceeds `max_sub_docs_per_row`. + #[serde(default)] + pub(crate) max_sub_docs_per_row_exceed_action: MaxSubDocsPerRowExceedAction, + /// Total memory limit in MiB for the build stage. /// /// This is split evenly across FTS workers at build time. By default Lance @@ -277,6 +336,8 @@ struct RawInvertedIndexParams { json_tokenizer_mode: Option, #[serde(default)] disable_cross_array_unnest: bool, + max_sub_docs_per_row: Option, + max_sub_docs_per_row_exceed_action: Option, #[serde(rename = "memory_limit", alias = "worker_memory_limit_mb")] memory_limit_mb: Option, #[serde(rename = "num_workers")] @@ -412,6 +473,10 @@ impl RawInvertedIndexParams { params.json_tokenizer_mode = Some(json_tokenizer_mode); } params.disable_cross_array_unnest = self.disable_cross_array_unnest; + params.max_sub_docs_per_row = self.max_sub_docs_per_row; + if let Some(action) = self.max_sub_docs_per_row_exceed_action { + params.max_sub_docs_per_row_exceed_action = action; + } params.memory_limit_mb = self.memory_limit_mb; params.num_workers = self.num_workers; params.format_version = self.format_version; @@ -452,6 +517,10 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { .filter(|mode| *mode == JsonTokenizerMode::FlattenedSubDocs) .map(|mode| mode.as_ref().to_string()), disable_cross_array_unnest: params.disable_cross_array_unnest, + max_sub_docs_per_row: params.max_sub_docs_per_row.map(|value| value as u64), + max_sub_docs_per_row_exceed_action: PbMaxSubDocsPerRowExceedAction::from( + params.max_sub_docs_per_row_exceed_action, + ) as i32, }) } } @@ -477,6 +546,10 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { .map(JsonTokenizerMode::from_str) .transpose()?, disable_cross_array_unnest: details.disable_cross_array_unnest, + max_sub_docs_per_row: details.max_sub_docs_per_row.map(|value| value as usize), + max_sub_docs_per_row_exceed_action: details + .max_sub_docs_per_row_exceed_action + .try_into()?, ..Self::default() }; params.document_granularity = details.document_granularity.try_into()?; @@ -535,6 +608,9 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { .map(JsonTokenizerMode::from_str) .transpose()?; params.disable_cross_array_unnest = details.disable_cross_array_unnest; + params.max_sub_docs_per_row = details.max_sub_docs_per_row.map(|value| value as usize); + params.max_sub_docs_per_row_exceed_action = + details.max_sub_docs_per_row_exceed_action.try_into()?; params.validate()?; Ok(params) } @@ -691,6 +767,8 @@ impl InvertedIndexParams { index_operators: false, json_tokenizer_mode: None, disable_cross_array_unnest: false, + max_sub_docs_per_row: None, + max_sub_docs_per_row_exceed_action: MaxSubDocsPerRowExceedAction::Fail, memory_limit_mb: None, num_workers: None, format_version: None, @@ -806,6 +884,40 @@ impl InvertedIndexParams { self } + /// Limit the number of flattened sub-documents emitted for one JSON row. + /// + /// If unset, the number of sub-documents is unlimited. + /// + /// # Examples + /// + /// ``` + /// use lance_index::scalar::{InvertedIndexParams, MaxSubDocsPerRowExceedAction}; + /// + /// let params = InvertedIndexParams::default() + /// .max_sub_docs_per_row(1024)? + /// .max_sub_docs_per_row_exceed_action(MaxSubDocsPerRowExceedAction::SkipRow); + /// assert!(params.build().is_ok()); + /// # Ok::<(), Box>(()) + /// ``` + pub fn max_sub_docs_per_row(mut self, max_sub_docs_per_row: usize) -> Result { + if max_sub_docs_per_row == 0 { + return Err(Error::invalid_input( + "max_sub_docs_per_row must be greater than zero".to_string(), + )); + } + self.max_sub_docs_per_row = Some(max_sub_docs_per_row); + Ok(self) + } + + /// Set the action taken when a JSON row exceeds `max_sub_docs_per_row`. + pub fn max_sub_docs_per_row_exceed_action( + mut self, + action: MaxSubDocsPerRowExceedAction, + ) -> Self { + self.max_sub_docs_per_row_exceed_action = action; + self + } + /// Set the lexical tokenizer implementation. /// /// Setting this to `"code"` selects the code analyzer defaults. Setting @@ -1102,7 +1214,11 @@ impl InvertedIndexParams { self.json_tokenizer_mode .unwrap_or(JsonTokenizerMode::SingleDocument), ) - .with_disable_cross_array_unnest(self.disable_cross_array_unnest), + .with_disable_cross_array_unnest(self.disable_cross_array_unnest) + .with_sub_doc_limit( + self.max_sub_docs_per_row, + self.max_sub_docs_per_row_exceed_action, + ), )), None => Ok(Box::new(TextTokenizer::new(tokenizer))), _ => Err(Error::invalid_input(format!( @@ -1129,6 +1245,11 @@ impl InvertedIndexParams { fn validate(&self) -> Result<()> { validate_block_size(self.block_size)?; + if self.max_sub_docs_per_row == Some(0) { + return Err(Error::invalid_input( + "max_sub_docs_per_row must be greater than zero".to_string(), + )); + } if self.base_tokenizer != "code" && (self.split_identifiers || self.split_on_numerics @@ -1216,7 +1337,10 @@ mod tests { use crate::pbold; use crate::pbold::inverted_index_details::DocumentGranularity as PbDocumentGranularity; - use super::{DocumentGranularity, InvertedIndexParams, InvertedListFormatVersion}; + use super::{ + DocumentGranularity, InvertedIndexParams, InvertedListFormatVersion, + MaxSubDocsPerRowExceedAction, + }; use lance_core::Error; use lance_tokenizer::{Language, TokenStream}; use rstest::rstest; @@ -1653,6 +1777,9 @@ mod tests { posting_format_version: None, json_tokenizer_mode: None, disable_cross_array_unnest: false, + max_sub_docs_per_row: None, + max_sub_docs_per_row_exceed_action: + pbold::inverted_index_details::MaxSubDocsPerRowExceedAction::Fail as i32, }; let params = InvertedIndexParams::try_from(&old_details).unwrap(); assert_eq!(params.block_size, 128); @@ -1676,6 +1803,22 @@ mod tests { ); } + #[test] + fn test_json_sub_doc_limit_details_conversion() { + let params = InvertedIndexParams::default() + .max_sub_docs_per_row(128) + .unwrap() + .max_sub_docs_per_row_exceed_action(MaxSubDocsPerRowExceedAction::SkipRow); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + let roundtrip = InvertedIndexParams::try_from(&details).unwrap(); + + assert_eq!(roundtrip.max_sub_docs_per_row, Some(128)); + assert_eq!( + roundtrip.max_sub_docs_per_row_exceed_action, + MaxSubDocsPerRowExceedAction::SkipRow + ); + } + #[rstest] #[case::block_size_128(128)] #[case::block_size_256(256)] diff --git a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs index 4672e0f4b62..6519732beb5 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs @@ -10,6 +10,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::str::FromStr; +use super::MaxSubDocsPerRowExceedAction; + /// Document type for full text search. #[derive(Debug, Clone)] pub enum DocType { @@ -179,6 +181,8 @@ pub struct JsonTokenizer { tokenizer: TextAnalyzer, mode: JsonTokenizerMode, disable_cross_array_unnest: bool, + max_sub_docs_per_row: Option, + max_sub_docs_per_row_exceed_action: MaxSubDocsPerRowExceedAction, } impl JsonTokenizer { @@ -187,6 +191,8 @@ impl JsonTokenizer { tokenizer, mode: JsonTokenizerMode::SingleDocument, disable_cross_array_unnest: false, + max_sub_docs_per_row: None, + max_sub_docs_per_row_exceed_action: MaxSubDocsPerRowExceedAction::Fail, } } @@ -203,6 +209,16 @@ impl JsonTokenizer { self.disable_cross_array_unnest = disable_cross_array_unnest; self } + + pub(crate) fn with_sub_doc_limit( + mut self, + max_sub_docs_per_row: Option, + exceed_action: MaxSubDocsPerRowExceedAction, + ) -> Self { + self.max_sub_docs_per_row = max_sub_docs_per_row; + self.max_sub_docs_per_row_exceed_action = exceed_action; + self + } } impl std::fmt::Debug for JsonTokenizer { @@ -213,6 +229,11 @@ impl std::fmt::Debug for JsonTokenizer { "disable_cross_array_unnest", &self.disable_cross_array_unnest, ) + .field("max_sub_docs_per_row", &self.max_sub_docs_per_row) + .field( + "max_sub_docs_per_row_exceed_action", + &self.max_sub_docs_per_row_exceed_action, + ) .finish() } } @@ -247,12 +268,29 @@ impl LanceTokenizer for JsonTokenizer { flatten_json(&value, "", &mut tokens, &mut position, &mut self.tokenizer); Ok(vec![tokens]) } - JsonTokenizerMode::FlattenedSubDocs => Ok(flatten_json_sub_docs( - &value, - "", - &mut self.tokenizer, - self.disable_cross_array_unnest, - )), + JsonTokenizerMode::FlattenedSubDocs => { + let sub_docs = flatten_json_sub_docs( + &value, + "", + &mut self.tokenizer, + self.disable_cross_array_unnest, + self.max_sub_docs_per_row, + ); + match sub_docs { + Ok(sub_docs) => Ok(sub_docs), + Err(max_sub_docs_per_row) => match self.max_sub_docs_per_row_exceed_action { + MaxSubDocsPerRowExceedAction::Fail => Err(Error::invalid_input(format!( + "JSON row exceeds max_sub_docs_per_row={max_sub_docs_per_row}; increase max_sub_docs_per_row or set disable_cross_array_unnest=true" + ))), + MaxSubDocsPerRowExceedAction::SkipRow => { + log::warn!( + "skipping JSON row that exceeds max_sub_docs_per_row={max_sub_docs_per_row}" + ); + Ok(Vec::new()) + } + }, + } + } } } @@ -440,24 +478,31 @@ fn flatten_json_sub_docs( prefix: &str, tokenizer: &mut TextAnalyzer, disable_cross_array_unnest: bool, -) -> Vec> { - flatten_json_sub_doc_terms(value, prefix, tokenizer, disable_cross_array_unnest) - .sub_docs - .into_iter() - .map(|sub_doc| { - sub_doc - .into_iter() - .enumerate() - .map(|(position, text)| Token { - offset_from: 0, - offset_to: 0, - position, - text, - position_length: 1, - }) - .collect() - }) - .collect() + max_sub_docs_per_row: Option, +) -> std::result::Result>, usize> { + Ok(flatten_json_sub_doc_terms( + value, + prefix, + tokenizer, + disable_cross_array_unnest, + max_sub_docs_per_row, + )? + .sub_docs + .into_iter() + .map(|sub_doc| { + sub_doc + .into_iter() + .enumerate() + .map(|(position, text)| Token { + offset_from: 0, + offset_to: 0, + position, + text, + position_length: 1, + }) + .collect() + }) + .collect()) } struct FlattenedJsonSubDocs { @@ -470,7 +515,8 @@ fn flatten_json_sub_doc_terms( prefix: &str, tokenizer: &mut TextAnalyzer, disable_cross_array_unnest: bool, -) -> FlattenedJsonSubDocs { + max_sub_docs_per_row: Option, +) -> std::result::Result { match value { Value::Object(map) => { let mut non_nested = Vec::new(); @@ -487,7 +533,8 @@ fn flatten_json_sub_doc_terms( &child_prefix, tokenizer, disable_cross_array_unnest, - ); + max_sub_docs_per_row, + )?; match child_terms.sub_docs.len() { 0 => {} 1 if !child_terms.has_array => { @@ -498,17 +545,17 @@ fn flatten_json_sub_doc_terms( } if nested.is_empty() { - return FlattenedJsonSubDocs { + return Ok(FlattenedJsonSubDocs { sub_docs: if non_nested.is_empty() { Vec::new() } else { vec![non_nested] }, has_array: false, - }; + }); } if nested.len() == 1 { - return FlattenedJsonSubDocs { + return Ok(FlattenedJsonSubDocs { sub_docs: nested .into_iter() .flatten() @@ -518,17 +565,17 @@ fn flatten_json_sub_doc_terms( }) .collect(), has_array: true, - }; + }); } let sub_docs = if disable_cross_array_unnest { - unnest_json_sub_docs(&nested, &non_nested) + unnest_json_sub_docs(&nested, &non_nested, max_sub_docs_per_row)? } else { - cross_join_json_sub_docs(&nested, &non_nested) + cross_join_json_sub_docs(&nested, &non_nested, max_sub_docs_per_row)? }; - FlattenedJsonSubDocs { + Ok(FlattenedJsonSubDocs { sub_docs, has_array: true, - } + }) } Value::Array(arr) => { let mut sub_docs = Vec::new(); @@ -539,16 +586,22 @@ fn flatten_json_sub_doc_terms( &child_prefix, tokenizer, disable_cross_array_unnest, - ); + max_sub_docs_per_row, + )?; for sub_doc in &mut child_terms.sub_docs { sub_doc.push(format!("{prefix}$idx,number,{array_index}")); } + if let Some(limit) = max_sub_docs_per_row + && sub_docs.len() + child_terms.sub_docs.len() > limit + { + return Err(limit); + } sub_docs.extend(child_terms.sub_docs); } - FlattenedJsonSubDocs { + Ok(FlattenedJsonSubDocs { sub_docs, has_array: true, - } + }) } Value::String(text) => { let mut token_texts = Vec::new(); @@ -556,46 +609,67 @@ fn flatten_json_sub_doc_terms( while let Some(token) = tokens.next() { token_texts.push(format!("{prefix},str,{}", token.text)); } - FlattenedJsonSubDocs { + Ok(FlattenedJsonSubDocs { sub_docs: if token_texts.is_empty() { Vec::new() } else { vec![token_texts] }, has_array: false, - } + }) } - Value::Null => FlattenedJsonSubDocs { + Value::Null => Ok(FlattenedJsonSubDocs { sub_docs: vec![vec![format!("{prefix},null,null")]], has_array: false, - }, - Value::Bool(value) => FlattenedJsonSubDocs { + }), + Value::Bool(value) => Ok(FlattenedJsonSubDocs { sub_docs: vec![vec![format!("{prefix},bool,{value}")]], has_array: false, - }, - Value::Number(value) => FlattenedJsonSubDocs { + }), + Value::Number(value) => Ok(FlattenedJsonSubDocs { sub_docs: vec![vec![format!("{prefix},number,{value}")]], has_array: false, - }, + }), } } fn cross_join_json_sub_docs( nested: &[Vec>], non_nested: &[String], -) -> Vec> { - let capacity = nested - .iter() - .map(|sub_docs| sub_docs.len()) - .product::(); + max_sub_docs_per_row: Option, +) -> std::result::Result>, usize> { + let capacity = if let Some(limit) = max_sub_docs_per_row { + let mut capacity = 1usize; + for sub_docs in nested { + capacity = capacity.saturating_mul(sub_docs.len()); + if capacity > limit { + return Err(limit); + } + } + capacity + } else { + nested + .iter() + .map(|sub_docs| sub_docs.len()) + .product::() + }; let mut results = Vec::with_capacity(capacity); let mut current = Vec::new(); cross_join_json_sub_docs_inner(nested, 0, non_nested, &mut current, &mut results); - results + Ok(results) } -fn unnest_json_sub_docs(nested: &[Vec>], non_nested: &[String]) -> Vec> { +fn unnest_json_sub_docs( + nested: &[Vec>], + non_nested: &[String], + max_sub_docs_per_row: Option, +) -> std::result::Result>, usize> { let capacity = nested.iter().map(|sub_docs| sub_docs.len()).sum::(); + if let Some(limit) = max_sub_docs_per_row + && capacity > limit + { + return Err(limit); + } let mut results = Vec::with_capacity(capacity); for sub_docs in nested { for child in sub_docs { @@ -604,7 +678,7 @@ fn unnest_json_sub_docs(nested: &[Vec>], non_nested: &[String]) -> V results.push(sub_doc); } } - results + Ok(results) } fn cross_join_json_sub_docs_inner( @@ -655,6 +729,7 @@ impl TokenStream for TTStream { #[cfg(test)] mod tests { + use crate::scalar::inverted::tokenizer::MaxSubDocsPerRowExceedAction; use crate::scalar::inverted::tokenizer::document_tokenizer::{ JsonTokenizer, JsonTokenizerMode, LanceTokenizer, flatten_json, flatten_json_sub_docs, flatten_triplet, @@ -828,6 +903,32 @@ mod tests { ); } + #[test] + fn test_max_sub_docs_per_row_actions() { + let json = r#"{"a":["x","y"],"b":["u","v"]}"#; + let tokenizer = || { + JsonTokenizer::new(TextAnalyzer::builder(SimpleTokenizer::default()).build()) + .with_mode(JsonTokenizerMode::FlattenedSubDocs) + }; + + let error = tokenizer() + .with_sub_doc_limit(Some(3), MaxSubDocsPerRowExceedAction::Fail) + .token_streams_for_doc(json) + .unwrap_err(); + assert!(error.to_string().contains("max_sub_docs_per_row=3")); + assert!( + error + .to_string() + .contains("disable_cross_array_unnest=true") + ); + + let sub_docs = tokenizer() + .with_sub_doc_limit(Some(3), MaxSubDocsPerRowExceedAction::SkipRow) + .token_streams_for_doc(json) + .unwrap(); + assert!(sub_docs.is_empty()); + } + fn assert_token(token: &Token, position: usize, text: &str) { assert_eq!( token.position, position, @@ -844,7 +945,8 @@ mod tests { let value: Value = serde_json::from_str(json).unwrap(); let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); let mut actual = - flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest) + flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest, None) + .unwrap() .into_iter() .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) .collect::>(); From ae2c627a03f93db34915d6b757df34f0e1b759aa Mon Sep 17 00:00:00 2001 From: wirybeaver Date: Thu, 17 Sep 2026 19:58:23 -0700 Subject: [PATCH 5/5] refactor(index): simplify flattened JSON sub-doc processing --- rust/lance-index/src/scalar/inverted.rs | 2 +- .../src/scalar/inverted/builder.rs | 19 +- .../src/scalar/inverted/index/flat_search.rs | 38 +- rust/lance-index/src/scalar/inverted/query.rs | 78 +++- .../inverted/tokenizer/document_tokenizer.rs | 366 +++++++----------- rust/lance/src/dataset/mem_wal/index/fts.rs | 10 +- rust/lance/src/io/exec/fts.rs | 133 +++---- 7 files changed, 319 insertions(+), 327 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 0ed34b82dca..c9f9cc55bb5 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -62,7 +62,7 @@ pub(crate) fn collapse_scored_rows( .total_cmp(left_score) .then_with(|| left_id.cmp(right_id)) }); - rows.truncate(rows.len().min(limit)); + rows.truncate(limit); rows } diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index b6e47b6be8c..c5501af87f5 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -206,12 +206,6 @@ impl InvertedIndexBuilder { } } - fn infer_lance_tokenizer(&mut self, doc_type: DocType) { - if self.params.lance_tokenizer.is_none() { - self.params.lance_tokenizer = Some(doc_type.as_ref().to_string()); - } - } - pub fn with_posting_tail_codec(mut self, posting_tail_codec: PostingTailCodec) -> Self { self.format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( posting_tail_codec, @@ -255,7 +249,9 @@ impl InvertedIndexBuilder { // infer lance_tokenizer based on document type let field = schema.column_with_name(doc_col).expect_ok()?.1; let doc_type = DocType::try_from(field)?; - self.infer_lance_tokenizer(doc_type); + self.params + .lance_tokenizer + .get_or_insert_with(|| doc_type.as_ref().to_string()); if self.params.lance_tokenizer.as_deref() == Some("json") && self.params.json_tokenizer_mode.is_none() { @@ -291,7 +287,9 @@ impl InvertedIndexBuilder { let field = schema.column_with_name(doc_col).expect_ok()?.1; let doc_type = DocType::try_from(field)?; - self.infer_lance_tokenizer(doc_type); + self.params + .lance_tokenizer + .get_or_insert_with(|| doc_type.as_ref().to_string()); let mut files = self .merge_existing_segments(dest_store, old_segments, old_data_filter.as_ref()) @@ -1594,10 +1592,9 @@ impl IndexWorker { } }; self.total_doc_length += doc.len(); - let with_position = self.has_position(); let sub_docs = self.tokenizer.token_streams_for_doc(doc.as_ref())?; for tokens in sub_docs { - self.process_tokenized_doc(row_id, tokens, with_position, doc_index) + self.process_tokenized_doc(row_id, tokens, doc_index) .await?; } @@ -1608,9 +1605,9 @@ impl IndexWorker { &mut self, row_id: u64, tokens: Vec, - with_position: bool, doc_index: &[u32], ) -> Result<()> { + let with_position = self.has_position(); let builder_was_empty = self.builder.docs.is_empty(); let old_temporary_memory_size = self.temporary_memory_size(); let old_token_memory_size = self.builder.tokens.memory_size() as u64; diff --git a/rust/lance-index/src/scalar/inverted/index/flat_search.rs b/rust/lance-index/src/scalar/inverted/index/flat_search.rs index d4f1f242f75..ef7cadb153b 100644 --- a/rust/lance-index/src/scalar/inverted/index/flat_search.rs +++ b/rust/lance-index/src/scalar/inverted/index/flat_search.rs @@ -69,6 +69,9 @@ pub(super) fn do_flat_full_text_search( let row_id_array = batch[ROW_ID].as_primitive::(); let doc_array = batch[doc_col].as_string::(); for i in 0..row_id_array.len() { + if doc_array.is_null(i) { + continue; + } let doc = doc_array.value(i); if document_matches_flat_query(doc, &mut tokenizer, &query_tokens, phrase_slop)? { results.push(row_id_array.value(i)); @@ -115,9 +118,14 @@ pub(super) fn do_flat_full_text_search_list( .join(" "); document_matches_flat_query(&document, &mut tokenizer, &query_tokens, phrase_slop)? } else { - iter_str_array(elements.as_ref()) - .flatten() - .any(|element| has_query_token(element, &mut tokenizer, &query_tokens)) + let mut matches = false; + for element in iter_str_array(elements.as_ref()).flatten() { + if document_matches_flat_query(element, &mut tokenizer, &query_tokens, None)? { + matches = true; + break; + } + } + matches }; if matches { results.push(row_id_array.value(i)); @@ -135,7 +143,12 @@ pub(super) fn document_matches_flat_query( phrase_slop: Option, ) -> Result { let Some(slop) = phrase_slop else { - return Ok(has_query_token(document, tokenizer, query_tokens)); + let operator = effective_json_query_operator( + tokenizer.json_tokenizer_mode(), + query_tokens, + Operator::Or, + ); + return document_matches_query(document, tokenizer, query_tokens, operator); }; for tokens in tokenizer.token_streams_for_doc(document)? { @@ -574,15 +587,14 @@ pub(super) fn initialize_scorer( } fn collapse_flattened_rows(batch: RecordBatch) -> Result { - let row_ids = batch[ROW_ID].as_primitive::().values().to_vec(); - let scores = batch[SCORE_COL] - .as_primitive::() - .values() - .to_vec(); - let (row_ids, scores): (Vec<_>, Vec<_>) = - collapse_scored_rows(row_ids.into_iter().zip(scores), usize::MAX) - .into_iter() - .unzip(); + let row_ids = batch[ROW_ID].as_primitive::().values(); + let scores = batch[SCORE_COL].as_primitive::().values(); + let (row_ids, scores): (Vec<_>, Vec<_>) = collapse_scored_rows( + row_ids.iter().copied().zip(scores.iter().copied()), + usize::MAX, + ) + .into_iter() + .unzip(); Ok(RecordBatch::try_new( FTS_SCHEMA.clone(), vec![ diff --git a/rust/lance-index/src/scalar/inverted/query.rs b/rust/lance-index/src/scalar/inverted/query.rs index f7c56c6a909..6f84fe5a1e1 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -5,6 +5,7 @@ use crate::scalar::inverted::DocumentGranularity; use crate::scalar::inverted::document_tokenizer::{DocType, JsonTokenizerMode}; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; use lance_core::{Error, Result}; +use roaring::RoaringBitmap; use serde::ser::SerializeMap; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; @@ -918,24 +919,63 @@ pub fn effective_json_query_operator( } } -pub fn has_query_token( +/// Match tokens within a single internal document, propagating tokenization errors. +/// +/// For JSON arrays, all required query positions must match the same sub-document. +/// Tokens sharing a query position are alternatives. +/// +/// ``` +/// use lance_index::scalar::InvertedIndexParams; +/// use lance_index::scalar::inverted::query::{ +/// collect_query_tokens, document_matches_query, Operator, +/// }; +/// +/// let mut tokenizer = InvertedIndexParams::default().build()?; +/// let tokens = collect_query_tokens("alpha beta", &mut tokenizer); +/// assert!(document_matches_query("alpha beta", &mut tokenizer, &tokens, Operator::And)?); +/// # Ok::<(), lance_core::Error>(()) +/// ``` +pub fn document_matches_query( text: &str, tokenizer: &mut Box, query_tokens: &Tokens, -) -> bool { - match tokenizer.token_streams_for_doc(text) { - Ok(sub_docs) => { - for tokens in sub_docs { - for token in tokens { - if query_tokens.contains(&token.text) { - return true; + operator: Operator, +) -> Result { + let sub_docs = tokenizer.token_streams_for_doc(text)?; + Ok(sub_docs.into_iter().any(|tokens| match operator { + Operator::Or => tokens + .iter() + .any(|token| query_tokens.contains(&token.text)), + Operator::And => { + let mut remaining_positions = (0..query_tokens.len()) + .map(|index| query_tokens.position(index)) + .collect::(); + if remaining_positions.is_empty() { + return false; + } + for token in tokens { + for index in 0..query_tokens.len() { + if token.text == query_tokens.get_token(index) { + remaining_positions.remove(query_tokens.position(index)); } } + if remaining_positions.is_empty() { + return true; + } } false } - Err(_) => false, - } + })) +} + +/// Return whether any query token matches, treating tokenization errors as no match. +#[deprecated(note = "use document_matches_query to propagate tokenization errors")] +pub fn has_query_token( + text: &str, + tokenizer: &mut Box, + query_tokens: &Tokens, +) -> bool { + document_matches_query(text, tokenizer, query_tokens, Operator::Or).unwrap_or(false) } fn fill_match_query_columns( @@ -1033,6 +1073,24 @@ pub fn fill_fts_query_column( #[cfg(test)] mod tests { use super::*; + use crate::scalar::InvertedIndexParams; + + #[test] + fn test_document_match_query_position_alternatives() { + let mut tokenizer = InvertedIndexParams::default().build().unwrap(); + let tokens = Tokens::with_positions( + vec!["red".to_string(), "blue".to_string(), "shoe".to_string()], + vec![0, 0, 1], + DocType::Text, + ); + assert!( + document_matches_query("blue shoe", &mut tokenizer, &tokens, Operator::And).unwrap() + ); + assert!( + !document_matches_query("red blue", &mut tokenizer, &tokens, Operator::And).unwrap() + ); + assert!(document_matches_query("red blue", &mut tokenizer, &tokens, Operator::Or).unwrap()); + } #[test] fn test_fuzzy_expansion_mode_and_unicode_auto_boundaries() { diff --git a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs index 6519732beb5..e8fc4bffba8 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs @@ -269,13 +269,12 @@ impl LanceTokenizer for JsonTokenizer { Ok(vec![tokens]) } JsonTokenizerMode::FlattenedSubDocs => { - let sub_docs = flatten_json_sub_docs( - &value, - "", - &mut self.tokenizer, - self.disable_cross_array_unnest, - self.max_sub_docs_per_row, - ); + let sub_docs = JsonFlattener { + tokenizer: &mut self.tokenizer, + disable_cross_array_unnest: self.disable_cross_array_unnest, + max_sub_docs_per_row: self.max_sub_docs_per_row, + } + .tokenize(&value); match sub_docs { Ok(sub_docs) => Ok(sub_docs), Err(max_sub_docs_per_row) => match self.max_sub_docs_per_row_exceed_action { @@ -325,12 +324,12 @@ fn flatten_triplet( let field = parts[0]; let v_type = parts[1]; let value = parts[2]; - let (field, mut index_tokens) = match mode { + let (field, index_tokens) = match mode { JsonTokenizerMode::SingleDocument => (field.to_string(), Vec::new()), JsonTokenizerMode::FlattenedSubDocs => normalize_flattened_json_path(field)?, }; - for index_token in index_tokens.drain(..) { + for index_token in index_tokens { token_vec.push(Token { offset_from: 0, offset_to: 0, @@ -379,18 +378,17 @@ fn flatten_triplet( fn normalize_flattened_json_path(path: &str) -> lance_core::Result<(String, Vec)> { let mut normalized = String::with_capacity(path.len()); let mut index_tokens = Vec::new(); - let mut chars = path.char_indices().peekable(); + let mut chars = path.chars(); - while let Some((_, ch)) = chars.next() { + while let Some(ch) = chars.next() { if ch != '[' { normalized.push(ch); continue; } - let index_path = normalized.clone(); let mut array_index = String::new(); let mut found_right_bracket = false; - for (_, bracket_ch) in chars.by_ref() { + for bracket_ch in chars.by_ref() { if bracket_ch == ']' { found_right_bracket = true; break; @@ -408,7 +406,7 @@ fn normalize_flattened_json_path(path: &str) -> lance_core::Result<(String, Vec< ))); } if array_index != "*" { - index_tokens.push(format!("{index_path}$idx,number,{array_index}")); + index_tokens.push(format!("{normalized}$idx,number,{array_index}")); } normalized.push('.'); } @@ -473,232 +471,159 @@ fn flatten_json( } } -fn flatten_json_sub_docs( - value: &Value, - prefix: &str, - tokenizer: &mut TextAnalyzer, - disable_cross_array_unnest: bool, - max_sub_docs_per_row: Option, -) -> std::result::Result>, usize> { - Ok(flatten_json_sub_doc_terms( - value, - prefix, - tokenizer, - disable_cross_array_unnest, - max_sub_docs_per_row, - )? - .sub_docs - .into_iter() - .map(|sub_doc| { - sub_doc - .into_iter() - .enumerate() - .map(|(position, text)| Token { - offset_from: 0, - offset_to: 0, - position, - text, - position_length: 1, - }) - .collect() - }) - .collect()) -} - struct FlattenedJsonSubDocs { sub_docs: Vec>, has_array: bool, } -fn flatten_json_sub_doc_terms( - value: &Value, - prefix: &str, - tokenizer: &mut TextAnalyzer, +struct JsonFlattener<'a> { + tokenizer: &'a mut TextAnalyzer, disable_cross_array_unnest: bool, max_sub_docs_per_row: Option, -) -> std::result::Result { - match value { - Value::Object(map) => { - let mut non_nested = Vec::new(); - let mut nested: Vec>> = Vec::new(); +} - for (key, child) in map { - let child_prefix = if prefix.is_empty() { - key.clone() +impl JsonFlattener<'_> { + fn tokenize(&mut self, value: &Value) -> Result>, usize> { + Ok(self + .flatten(value, "")? + .sub_docs + .into_iter() + .map(|sub_doc| { + sub_doc + .into_iter() + .enumerate() + .map(|(position, text)| Token { + offset_from: 0, + offset_to: 0, + position, + text, + position_length: 1, + }) + .collect() + }) + .collect()) + } + + fn check_count(&self, count: Option) -> Result { + let count = count.ok_or(self.max_sub_docs_per_row.unwrap_or(usize::MAX))?; + if let Some(limit) = self.max_sub_docs_per_row + && count > limit + { + return Err(limit); + } + Ok(count) + } + + fn flatten(&mut self, value: &Value, prefix: &str) -> Result { + match value { + Value::Object(map) => { + let mut scalar_terms = Vec::new(); + let mut array_groups = Vec::new(); + for (key, child) in map { + let child_prefix = if prefix.is_empty() { + key.clone() + } else { + format!("{prefix}.{key}") + }; + let child = self.flatten(child, &child_prefix)?; + if child.has_array { + if !child.sub_docs.is_empty() { + array_groups.push(child.sub_docs); + } + } else { + scalar_terms.extend(child.sub_docs.into_iter().flatten()); + } + } + if array_groups.is_empty() { + return Ok(FlattenedJsonSubDocs { + sub_docs: if scalar_terms.is_empty() { + Vec::new() + } else { + vec![scalar_terms] + }, + has_array: false, + }); + } + let mut sub_docs = if self.disable_cross_array_unnest || array_groups.len() == 1 { + let capacity = array_groups.iter().try_fold(0usize, |count, group| { + self.check_count(count.checked_add(group.len())) + })?; + let mut sub_docs = Vec::with_capacity(capacity); + sub_docs.extend(array_groups.into_iter().flatten()); + sub_docs } else { - format!("{prefix}.{key}") + let capacity = array_groups.iter().try_fold(1usize, |count, group| { + self.check_count(count.checked_mul(group.len())) + })?; + let mut sub_docs = Vec::with_capacity(capacity); + cross_join_json_sub_docs(&array_groups, &mut Vec::new(), &mut sub_docs); + sub_docs }; - let child_terms = flatten_json_sub_doc_terms( - child, - &child_prefix, - tokenizer, - disable_cross_array_unnest, - max_sub_docs_per_row, - )?; - match child_terms.sub_docs.len() { - 0 => {} - 1 if !child_terms.has_array => { - non_nested.extend(child_terms.sub_docs.into_iter().flatten()) + for sub_doc in &mut sub_docs { + sub_doc.extend(scalar_terms.iter().cloned()); + } + Ok(FlattenedJsonSubDocs { + sub_docs, + has_array: true, + }) + } + Value::Array(values) => { + let mut sub_docs = Vec::new(); + let child_prefix = format!("{prefix}."); + for (array_index, child) in values.iter().enumerate() { + let mut child = self.flatten(child, &child_prefix)?; + self.check_count(sub_docs.len().checked_add(child.sub_docs.len()))?; + for sub_doc in &mut child.sub_docs { + sub_doc.push(format!("{prefix}$idx,number,{array_index}")); } - _ => nested.push(child_terms.sub_docs), + sub_docs.extend(child.sub_docs); } + Ok(FlattenedJsonSubDocs { + sub_docs, + has_array: true, + }) } - - if nested.is_empty() { - return Ok(FlattenedJsonSubDocs { - sub_docs: if non_nested.is_empty() { + _ => { + let terms = match value { + Value::String(text) => { + let mut terms = Vec::new(); + let mut stream = self.tokenizer.token_stream(text); + while let Some(token) = stream.next() { + terms.push(format!("{prefix},str,{}", token.text)); + } + terms + } + Value::Null => vec![format!("{prefix},null,null")], + Value::Bool(value) => vec![format!("{prefix},bool,{value}")], + Value::Number(value) => vec![format!("{prefix},number,{value}")], + _ => unreachable!(), + }; + Ok(FlattenedJsonSubDocs { + sub_docs: if terms.is_empty() { Vec::new() } else { - vec![non_nested] + vec![terms] }, has_array: false, - }); - } - if nested.len() == 1 { - return Ok(FlattenedJsonSubDocs { - sub_docs: nested - .into_iter() - .flatten() - .map(|mut sub_doc| { - sub_doc.extend(non_nested.iter().cloned()); - sub_doc - }) - .collect(), - has_array: true, - }); - } - let sub_docs = if disable_cross_array_unnest { - unnest_json_sub_docs(&nested, &non_nested, max_sub_docs_per_row)? - } else { - cross_join_json_sub_docs(&nested, &non_nested, max_sub_docs_per_row)? - }; - Ok(FlattenedJsonSubDocs { - sub_docs, - has_array: true, - }) - } - Value::Array(arr) => { - let mut sub_docs = Vec::new(); - let child_prefix = format!("{prefix}."); - for (array_index, child) in arr.iter().enumerate() { - let mut child_terms = flatten_json_sub_doc_terms( - child, - &child_prefix, - tokenizer, - disable_cross_array_unnest, - max_sub_docs_per_row, - )?; - for sub_doc in &mut child_terms.sub_docs { - sub_doc.push(format!("{prefix}$idx,number,{array_index}")); - } - if let Some(limit) = max_sub_docs_per_row - && sub_docs.len() + child_terms.sub_docs.len() > limit - { - return Err(limit); - } - sub_docs.extend(child_terms.sub_docs); - } - Ok(FlattenedJsonSubDocs { - sub_docs, - has_array: true, - }) - } - Value::String(text) => { - let mut token_texts = Vec::new(); - let mut tokens = tokenizer.token_stream(text); - while let Some(token) = tokens.next() { - token_texts.push(format!("{prefix},str,{}", token.text)); + }) } - Ok(FlattenedJsonSubDocs { - sub_docs: if token_texts.is_empty() { - Vec::new() - } else { - vec![token_texts] - }, - has_array: false, - }) } - Value::Null => Ok(FlattenedJsonSubDocs { - sub_docs: vec![vec![format!("{prefix},null,null")]], - has_array: false, - }), - Value::Bool(value) => Ok(FlattenedJsonSubDocs { - sub_docs: vec![vec![format!("{prefix},bool,{value}")]], - has_array: false, - }), - Value::Number(value) => Ok(FlattenedJsonSubDocs { - sub_docs: vec![vec![format!("{prefix},number,{value}")]], - has_array: false, - }), } } fn cross_join_json_sub_docs( - nested: &[Vec>], - non_nested: &[String], - max_sub_docs_per_row: Option, -) -> std::result::Result>, usize> { - let capacity = if let Some(limit) = max_sub_docs_per_row { - let mut capacity = 1usize; - for sub_docs in nested { - capacity = capacity.saturating_mul(sub_docs.len()); - if capacity > limit { - return Err(limit); - } - } - capacity - } else { - nested - .iter() - .map(|sub_docs| sub_docs.len()) - .product::() - }; - let mut results = Vec::with_capacity(capacity); - let mut current = Vec::new(); - cross_join_json_sub_docs_inner(nested, 0, non_nested, &mut current, &mut results); - Ok(results) -} - -fn unnest_json_sub_docs( - nested: &[Vec>], - non_nested: &[String], - max_sub_docs_per_row: Option, -) -> std::result::Result>, usize> { - let capacity = nested.iter().map(|sub_docs| sub_docs.len()).sum::(); - if let Some(limit) = max_sub_docs_per_row - && capacity > limit - { - return Err(limit); - } - let mut results = Vec::with_capacity(capacity); - for sub_docs in nested { - for child in sub_docs { - let mut sub_doc = child.clone(); - sub_doc.extend(non_nested.iter().cloned()); - results.push(sub_doc); - } - } - Ok(results) -} - -fn cross_join_json_sub_docs_inner( - nested: &[Vec>], - nested_index: usize, - non_nested: &[String], + groups: &[Vec>], current: &mut Vec, results: &mut Vec>, ) { - if nested_index == nested.len() { - let mut sub_doc = current.clone(); - sub_doc.extend(non_nested.iter().cloned()); - results.push(sub_doc); + let Some((group, remaining)) = groups.split_first() else { + results.push(current.clone()); return; - } - - for child in &nested[nested_index] { + }; + for child in group { let old_len = current.len(); current.extend(child.iter().cloned()); - cross_join_json_sub_docs_inner(nested, nested_index + 1, non_nested, current, results); + cross_join_json_sub_docs(remaining, current, results); current.truncate(old_len); } } @@ -731,7 +656,7 @@ impl TokenStream for TTStream { mod tests { use crate::scalar::inverted::tokenizer::MaxSubDocsPerRowExceedAction; use crate::scalar::inverted::tokenizer::document_tokenizer::{ - JsonTokenizer, JsonTokenizerMode, LanceTokenizer, flatten_json, flatten_json_sub_docs, + JsonFlattener, JsonTokenizer, JsonTokenizerMode, LanceTokenizer, flatten_json, flatten_triplet, }; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer, Token}; @@ -915,6 +840,7 @@ mod tests { .with_sub_doc_limit(Some(3), MaxSubDocsPerRowExceedAction::Fail) .token_streams_for_doc(json) .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); assert!(error.to_string().contains("max_sub_docs_per_row=3")); assert!( error @@ -944,12 +870,16 @@ mod tests { fn assert_sub_docs(json: &str, disable_cross_array_unnest: bool, expected: &[&str]) { let value: Value = serde_json::from_str(json).unwrap(); let mut tokenizer = TextAnalyzer::builder(SimpleTokenizer::default()).build(); - let mut actual = - flatten_json_sub_docs(&value, "", &mut tokenizer, disable_cross_array_unnest, None) - .unwrap() - .into_iter() - .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) - .collect::>(); + let mut actual = JsonFlattener { + tokenizer: &mut tokenizer, + disable_cross_array_unnest, + max_sub_docs_per_row: None, + } + .tokenize(&value) + .unwrap() + .into_iter() + .map(|tokens| sorted_tokens(tokens.into_iter().map(|token| token.text))) + .collect::>(); actual.sort(); let mut expected = expected .iter() diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 1fe5b128df2..f24a66b84d6 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -2579,7 +2579,15 @@ impl FtsMemIndex { }) .or_insert(entry); } - rows.into_values().collect() + let mut rows = rows.into_values().collect::>(); + rows.sort_unstable_by(|left, right| { + right + .score + .total_cmp(&left.score) + .then_with(|| left.row_position.cmp(&right.row_position)) + .then_with(|| left.doc_index.cmp(&right.doc_index)) + }); + rows } // ------------------------------------------------------------------ diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index b36e28f375b..8bd0bbdd711 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -56,7 +56,8 @@ use lance_index::scalar::inverted::document_tokenizer::{ }; use lance_index::scalar::inverted::query::{ BoostQuery, FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, Operator, PhraseQuery, Tokens, - collect_query_tokens, effective_json_query_operator, uses_fuzzy_expansion, + collect_query_tokens, document_matches_query, effective_json_query_operator, + uses_fuzzy_expansion, }; use lance_index::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_index::scalar::inverted::{ @@ -3040,41 +3041,6 @@ struct FlatMatchFilterStreamOptions { metrics_set: ExecutionPlanMetricsSet, } -fn document_matches_query( - text: &str, - tokenizer: &mut Box, - query_tokens: &Tokens, - operator: Operator, -) -> bool { - let Ok(sub_docs) = tokenizer.token_streams_for_doc(text) else { - return false; - }; - sub_docs.into_iter().any(|tokens| match operator { - Operator::Or => tokens - .iter() - .any(|token| query_tokens.contains(&token.text)), - Operator::And => { - let mut remaining_positions = (0..query_tokens.len()) - .map(|index| query_tokens.position(index)) - .collect::>(); - if remaining_positions.is_empty() { - return false; - } - for token in tokens { - for index in 0..query_tokens.len() { - if token.text == query_tokens.get_token(index) { - remaining_positions.remove(&query_tokens.position(index)); - } - } - if remaining_positions.is_empty() { - return true; - } - } - false - } - }) -} - impl DisplayAs for FlatMatchFilterExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { @@ -3230,7 +3196,7 @@ impl FlatMatchFilterExec { tokenizer: &mut Box, query_tokens: &Tokens, operator: Operator, - ) -> BooleanArray { + ) -> Result { let text_col = text_col.as_string::(); let mut predicate = BooleanBuilder::with_capacity(text_col.len()); for idx in 0..text_col.len() { @@ -3241,10 +3207,10 @@ impl FlatMatchFilterExec { tokenizer, query_tokens, operator, - ), + )?, ); } - predicate.finish() + Ok(predicate.finish()) } async fn build_filter_stream( @@ -3335,7 +3301,7 @@ impl FlatMatchFilterExec { &mut tokenizer, &query_tokens, query_operator, - ) { + )? { matches[document.row_index] = true; } } @@ -3368,7 +3334,7 @@ impl FlatMatchFilterExec { column, ))); } - }; + }?; Ok(arrow::compute::filter_record_batch(&batch, &predicate)?) } }); @@ -4843,13 +4809,15 @@ mod tests { use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::query::{ BooleanQuery, BoostQuery, FtsQuery, FtsSearchParams, MatchQuery, Occur, Operator, - PhraseQuery, collect_query_tokens, has_query_token, + PhraseQuery, collect_query_tokens, document_matches_query, }; use lance_index::scalar::inverted::{ DocumentGranularity, FTS_SCHEMA, InvertedIndex, Language, SCORE_COL, - build_global_bm25_scorer, prepare_bm25_query, + build_global_bm25_scorer, flat_full_text_search, prepare_bm25_query, + }; + use lance_index::scalar::{ + FullTextSearchQuery, InvertedIndexParams, MaxSubDocsPerRowExceedAction, }; - use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams}; use lance_index::{IndexCriteria, IndexType}; use lance_table::format::IndexMetadata; use uuid::Uuid; @@ -4867,8 +4835,8 @@ mod tests { BoolSlot, BoostQueryExec, CompoundQueryExec, CrossColumnCompoundQueryExec, FTS_SEGMENT_BIND_DURATION_METRIC, FlatMatchFilterExec, FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec, WAND_TIE_COMPLETION_BUDGET, WandExactnessCertificate, - build_boolean_query_children, classify_wand_exactness_certificate, default_text_tokenizer, - open_fts_segments, tokenizer_for_match_query, + build_boolean_query_children, classify_wand_exactness_certificate, open_fts_segments, + tokenizer_for_match_query, }; use crate::io::exec::utils::IndexMetrics; use datafusion::physical_plan::empty::EmptyExec; @@ -5161,30 +5129,46 @@ mod tests { } #[test] - fn document_match_filter_respects_document_boundary() { - let mut tokenizer = default_text_tokenizer(); - let query_tokens = collect_query_tokens("alpha", &mut tokenizer); - assert!(super::document_matches_query( - "alpha beta", - &mut tokenizer, - &query_tokens, - Operator::Or, - )); - - let mut tokenizer = default_text_tokenizer(); - let query_tokens = collect_query_tokens("alpha beta", &mut tokenizer); - assert!(!super::document_matches_query( - "alpha", - &mut tokenizer, - &query_tokens, - Operator::And, - )); - assert!(super::document_matches_query( - "alpha beta", - &mut tokenizer, - &query_tokens, - Operator::And, - )); + fn json_flat_filters_propagate_sub_doc_limit() -> lance_core::Result<()> { + let params: InvertedIndexParams = serde_json::from_value(serde_json::json!({ + "lance_tokenizer": "json", + "json_tokenizer_mode": "flattened_sub_docs", + "max_sub_docs_per_row": 2 + }))?; + let batch = arrow_array::record_batch!( + ("_rowid", UInt64, [0, 1, 2]), + ( + "json", + Utf8, + [Some(r#"{"a":["x","x","x"]}"#), Some(r#"{"a":["x"]}"#), None] + ) + )?; + let text = "a[*],str,x"; + let query = collect_query_tokens(text, &mut params.build()?); + let filter = |params: &InvertedIndexParams| { + FlatMatchFilterExec::find_matches::( + batch["json"].as_ref(), + &mut params.build()?, + &query, + Operator::Or, + ) + }; + let scan = |params: &InvertedIndexParams| { + flat_full_text_search(&[&batch], "json", text, Some(params.build()?)) + }; + for result in [filter(¶ms).map(|_| ()), scan(¶ms).map(|_| ())] { + let error = result.unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains("max_sub_docs_per_row=2")); + } + let params = + params.max_sub_docs_per_row_exceed_action(MaxSubDocsPerRowExceedAction::SkipRow); + assert_eq!( + filter(¶ms)?, + arrow_array::BooleanArray::from(vec![false, true, false]) + ); + assert_eq!(scan(¶ms)?, vec![1]); + Ok(()) } #[tokio::test] @@ -5318,7 +5302,8 @@ mod tests { &mut tokenizer, &query_tokens, Operator::Or, - ); + ) + .unwrap(); assert_eq!(result.len(), 3); assert!(result.value(0), "expected match in 'hello world'"); @@ -5388,9 +5373,11 @@ mod tests { ) .await .unwrap(); - assert!(has_query_token("hello", &mut tokenizer, &query_tokens)); assert!( - !has_query_token("HELLO", &mut tokenizer, &query_tokens), + document_matches_query("hello", &mut tokenizer, &query_tokens, Operator::Or).unwrap() + ); + assert!( + !document_matches_query("HELLO", &mut tokenizer, &query_tokens, Operator::Or).unwrap(), "legacy FTS indices should continue using on-disk tokenizer params" ); }