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..bf2ec0ad534 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,9 @@ public static final class Builder { private Integer maxNgramLength; 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; @@ -304,6 +307,53 @@ 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; + } + + /** + * 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. * @@ -501,6 +551,16 @@ public ScalarIndexParams build() { if (blockSize != null) { params.put("block_size", blockSize); } + 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 59ef98db9e1..81a68d2fc0c 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,29 @@ 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 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 01be98faeb0..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. @@ -115,4 +120,18 @@ 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; + /* 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 6ae17bb1c00..f65825aa1aa 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3720,6 +3720,20 @@ 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. + 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 79762a9539f..6db248f19bc 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -6032,7 +6032,14 @@ def test_json_inverted_match_query(tmp_path): stem=True, 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( @@ -6048,7 +6055,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..1373e71b93f 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2649,6 +2649,9 @@ impl Dataset { "split_on_numerics", "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", @@ -2682,6 +2685,24 @@ impl Dataset { .block_size(block_size.extract()?) .map_err(|e| PyValueError::new_err(e.to_string()))?; } + 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(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.rs b/rust/lance-index/src/scalar/inverted.rs index c3070d4d61c..c9f9cc55bb5 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(limit); + rows +} + /// Canonical token vocabulary and BM25 statistics for one indexed query leaf. /// /// Keeping these values together prevents a search path from expanding one @@ -428,11 +454,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..c5501af87f5 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; @@ -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; @@ -231,6 +232,10 @@ impl InvertedIndexBuilder { self } + pub(crate) fn params(&self) -> &InvertedIndexParams { + &self.params + } + pub async fn update( &mut self, new_data: SendableRecordBatchStream, @@ -242,11 +247,15 @@ 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.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() + { + self.params.json_tokenizer_mode = Some(JsonTokenizerMode::FlattenedSubDocs); } let new_data = document_input(new_data, doc_col)?; @@ -276,11 +285,11 @@ 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.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()) @@ -1575,6 +1584,28 @@ impl IndexWorker { row_id: u64, document: DocumentSource<'_>, doc_index: &[u32], + ) -> Result<()> { + let doc = match document { + DocumentSource::Text(doc) => Cow::Borrowed(doc), + DocumentSource::StringList(elements) => { + Cow::Owned(Self::materialize_string_list(elements)) + } + }; + self.total_doc_length += doc.len(); + let sub_docs = self.tokenizer.token_streams_for_doc(doc.as_ref())?; + for tokens in sub_docs { + self.process_tokenized_doc(row_id, tokens, doc_index) + .await?; + } + + Ok(()) + } + + async fn process_tokenized_doc( + &mut self, + row_id: u64, + tokens: Vec, + doc_index: &[u32], ) -> Result<()> { let with_position = self.has_position(); let builder_was_empty = self.builder.docs.is_empty(); @@ -1582,100 +1613,58 @@ impl IndexWorker { 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 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()); + 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 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; - } - }; - - 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 +1717,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 +1781,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..f09608ca8b2 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -16,8 +16,8 @@ use lance_select::RowAddrMask; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; use super::{ - InvertedIndex, PreparedBm25Query, - document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}, + InvertedIndex, PreparedBm25Query, collapse_scored_rows, + document_tokenizer::{DocType, JsonTokenizer, JsonTokenizerMode, LanceTokenizer}, documents::{ CachedRowAddressOrder, DocId, DocLengths, DocVisibility, OrderedRowAddressProjection, PartitionDocuments, ResidentAddressProjection, RowAddressProjectionOrderError, @@ -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, @@ -3791,9 +3792,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).with_mode( + index_tokenizer + .json_tokenizer_mode() + .unwrap_or(JsonTokenizerMode::SingleDocument), + ), + ) as Box, } } else { index.tokenizer() @@ -3832,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( @@ -3848,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)) @@ -4353,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, @@ -4368,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 = @@ -4429,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.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..ef7cadb153b 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}") @@ -68,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)); @@ -114,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)); @@ -134,31 +143,36 @@ 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); }; - 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"; @@ -299,87 +313,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 +416,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 +431,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 +483,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 +500,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 +586,24 @@ pub(super) fn initialize_scorer( MemBM25Scorer::new(total_tokens, num_docs, token_counts_map) } +fn collapse_flattened_rows(batch: RecordBatch) -> Result { + 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![ + 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, @@ -876,6 +893,10 @@ 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 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 @@ -937,7 +958,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, @@ -946,6 +967,9 @@ pub async fn flat_bm25_search_stream_with_options_and_scorer( boost, phrase_slop, )?; + if should_deduplicate_rows { + scores = collapse_flattened_rows(scores)?; + } // 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/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 78a533d7acb..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"; @@ -587,22 +588,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 +628,29 @@ 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(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) } } @@ -1167,7 +1195,6 @@ impl InvertedIndex { Ok(resolved_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..6f84fe5a1e1 100644 --- a/rust/lance-index/src/scalar/inverted/query.rs +++ b/rust/lance-index/src/scalar/inverted/query.rs @@ -2,9 +2,10 @@ // 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 roaring::RoaringBitmap; use serde::ser::SerializeMap; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; @@ -899,18 +900,82 @@ 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 + } +} + +/// 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, + 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 + } + })) +} + +/// 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 { - let mut stream = tokenizer.token_stream_for_doc(text); - while let Some(token) = stream.next() { - if query_tokens.contains(&token.text) { - return true; - } - } - false + document_matches_query(text, tokenizer, query_tokens, Operator::Or).unwrap_or(false) } fn fill_match_query_columns( @@ -1008,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.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index dcc5b5c3ac7..9f14bdf53b0 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; @@ -20,8 +20,9 @@ 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, LanceTokenizer, TextTokenizer, + JsonTokenizer, JsonTokenizerMode, LanceTokenizer, TextTokenizer, }; use crate::scalar::inverted::{ InvertedListFormatVersion, default_fts_format_version_for_block_size, @@ -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 { @@ -196,6 +245,29 @@ 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, + + /// 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 @@ -261,6 +333,11 @@ struct RawInvertedIndexParams { split_on_numerics: Option, preserve_original: Option, index_operators: Option, + 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")] @@ -392,6 +469,14 @@ 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.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; @@ -427,6 +512,15 @@ 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, + 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, }) } } @@ -446,6 +540,16 @@ 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, + 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()?; @@ -498,6 +602,15 @@ 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.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) } @@ -652,6 +765,10 @@ impl InvertedIndexParams { split_on_numerics: false, preserve_original: false, 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, @@ -749,6 +866,58 @@ impl InvertedIndexParams { self.document_granularity } + /// 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 + } + + /// 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 @@ -1039,7 +1208,18 @@ 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) + .with_mode( + self.json_tokenizer_mode + .unwrap_or(JsonTokenizerMode::SingleDocument), + ) + .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!( "unknown lance tokenizer {}", @@ -1065,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 @@ -1152,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; @@ -1587,6 +1775,11 @@ mod tests { code_config: None, document_granularity: PbDocumentGranularity::Row as i32, 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); @@ -1610,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 62dd7b1aa3b..e8fc4bffba8 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,13 @@ 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; + +use super::MaxSubDocsPerRowExceedAction; /// Document type for full text search. #[derive(Debug, Clone)] @@ -14,6 +19,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 +116,23 @@ 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 + } } impl Clone for Box { @@ -128,39 +179,120 @@ impl LanceTokenizer for TextTokenizer { #[derive(Clone)] 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 { pub fn new(tokenizer: TextAnalyzer) -> Self { - Self { tokenizer } + Self { + tokenizer, + mode: JsonTokenizerMode::SingleDocument, + disable_cross_array_unnest: false, + max_sub_docs_per_row: None, + max_sub_docs_per_row_exceed_action: MaxSubDocsPerRowExceedAction::Fail, + } + } + + #[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 + } + + 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 { 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, + ) + .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() } } 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 => { + 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 { + 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()) + } + }, + } + } + } + } + fn box_clone(&self) -> Box { Box::new(self.clone()) } @@ -168,9 +300,17 @@ impl LanceTokenizer for JsonTokenizer { fn doc_type(&self) -> DocType { DocType::Json } + + fn json_tokenizer_mode(&self) -> Option { + Some(self.mode) + } } -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 +324,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, index_tokens) = match mode { + JsonTokenizerMode::SingleDocument => (field.to_string(), Vec::new()), + JsonTokenizerMode::FlattenedSubDocs => normalize_flattened_json_path(field)?, + }; + + for index_token in index_tokens { + 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 +375,45 @@ 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.chars(); + + while let Some(ch) = chars.next() { + if ch != '[' { + normalized.push(ch); + continue; + } + + 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!("{normalized}$idx,number,{array_index}")); + } + normalized.push('.'); + } + + Ok((normalized, index_tokens)) +} + fn flatten_json( value: &Value, prefix: &str, @@ -277,6 +471,163 @@ fn flatten_json( } } +struct FlattenedJsonSubDocs { + sub_docs: Vec>, + has_array: bool, +} + +struct JsonFlattener<'a> { + tokenizer: &'a mut TextAnalyzer, + disable_cross_array_unnest: bool, + max_sub_docs_per_row: Option, +} + +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 { + 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 + }; + 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}")); + } + sub_docs.extend(child.sub_docs); + } + Ok(FlattenedJsonSubDocs { + sub_docs, + has_array: true, + }) + } + _ => { + 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![terms] + }, + has_array: false, + }) + } + } + } +} + +fn cross_join_json_sub_docs( + groups: &[Vec>], + current: &mut Vec, + results: &mut Vec>, +) { + let Some((group, remaining)) = groups.split_first() else { + results.push(current.clone()); + return; + }; + for child in group { + let old_len = current.len(); + current.extend(child.iter().cloned()); + cross_join_json_sub_docs(remaining, current, results); + current.truncate(old_len); + } +} + struct TTStream { tokens: Vec, index: usize, @@ -303,8 +654,10 @@ impl TokenStream for TTStream { #[cfg(test)] mod tests { + use crate::scalar::inverted::tokenizer::MaxSubDocsPerRowExceedAction; use crate::scalar::inverted::tokenizer::document_tokenizer::{ - JsonTokenizer, LanceTokenizer, flatten_json, flatten_triplet, + JsonFlattener, JsonTokenizer, JsonTokenizerMode, LanceTokenizer, flatten_json, + flatten_triplet, }; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer, Token}; use serde_json::Value; @@ -368,7 +721,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 +733,128 @@ mod tests { assert_token(&tokens[5], 5, "e,number,1.0"); } + #[test] + 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_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(); + 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_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_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_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", + ], + ); + } + + #[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!(matches!(error, lance_core::Error::InvalidInput { .. })); + 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, @@ -390,4 +866,40 @@ mod tests { "expected text {text} but {token:?}" ); } + + 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 = 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() + .map(|sub_doc| sorted_tokens(sub_doc.split(';'))) + .collect::>(); + expected.sort(); + assert_eq!(actual, expected); + } + + fn sorted_tokens(tokens: impl IntoIterator>) -> Vec { + let mut tokens = tokens.into_iter().map(Into::into).collect::>(); + tokens.sort(); + tokens + } + + 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/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 50ab8cc3b36..f24a66b84d6 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,47 @@ 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); + } + 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 + } + // ------------------------------------------------------------------ // Flush to Lance inverted index format // ------------------------------------------------------------------ @@ -2578,10 +2656,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 +2844,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 +2860,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 2d7c48c72a1..5a6f5517027 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -6457,6 +6457,89 @@ async fn prepare_json_dataset() -> (Dataset, String) { (dataset, json_col) } +#[tokio::test] +async fn test_json_inverted_flattened_sub_docs() { + let json_col = "json_field".to_string(); + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + 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, + ], + ) + .unwrap() + }; + + 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 stream = RecordBatchIterator::new(vec![Ok(indexed_batch)], schema.clone()); + let mut dataset = Dataset::write(stream, "memory://", 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 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], + ), + ]; + 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] async fn test_json_inverted_fuzziness_query() { let (mut dataset, json_col) = prepare_json_dataset().await; @@ -6817,7 +6900,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 +6942,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..8bd0bbdd711 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -51,10 +51,13 @@ 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, + 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::{ @@ -2052,7 +2055,13 @@ 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).with_mode(mode)) + } } } @@ -2928,12 +2937,17 @@ impl ExecutionPlan for MatchQueryExec { )))?; let mut tokenizer = tokenizer_for_match_query(first_index, query.fuzziness); 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 { query: prepared_query, params: Arc::new(params), - operator: query.operator, + operator, }) } else { let base_scorer = match (preset_base_scorer, shared_scorer) { @@ -2954,7 +2968,7 @@ impl ExecutionPlan for MatchQueryExec { &indices, tokens, params, - query.operator, + operator, metrics.as_ref(), base_scorer, ) @@ -3027,37 +3041,6 @@ struct FlatMatchFilterStreamOptions { metrics_set: ExecutionPlanMetricsSet, } -fn document_matches_query( - text: &str, - tokenizer: &mut Box, - query_tokens: &Tokens, - operator: Operator, -) -> bool { - match operator { - Operator::Or => has_query_token(text, tokenizer, query_tokens), - Operator::And => { - let mut remaining_positions = (0..query_tokens.len()) - .map(|index| query_tokens.position(index)) - .collect::>(); - if remaining_positions.is_empty() { - return false; - } - let mut stream = tokenizer.token_stream_for_doc(text); - while let Some(token) = stream.next() { - 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 { @@ -3213,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() { @@ -3224,10 +3207,10 @@ impl FlatMatchFilterExec { tokenizer, query_tokens, operator, - ), + )?, ); } - predicate.finish() + Ok(predicate.finish()) } async fn build_filter_stream( @@ -3289,6 +3272,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); @@ -3299,7 +3287,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(); @@ -3314,7 +3301,7 @@ impl FlatMatchFilterExec { &mut tokenizer, &query_tokens, query_operator, - ) { + )? { matches[document.row_index] = true; } } @@ -3347,7 +3334,7 @@ impl FlatMatchFilterExec { column, ))); } - }; + }?; Ok(arrow::compute::filter_record_batch(&batch, &predicate)?) } }); @@ -4822,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; @@ -4846,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; @@ -5140,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] @@ -5297,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'"); @@ -5367,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" ); }