diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index ede00ad43ee..e3c2569f774 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -13,12 +13,12 @@ use crate::vector::graph::OrderedFloat; use crate::{progress::IndexBuildProgress, progress::noop_progress}; use arrow::array::AsArray; use arrow::datatypes; -use arrow_array::{Array, BinaryArray, RecordBatch, UInt64Array}; +use arrow_array::{Array, BinaryArray, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use bytes::Bytes; -use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion::execution::SendableRecordBatchStream; use fst::Streamer; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{StreamExt, TryStreamExt}; use lance_arrow::json::JSON_EXT_NAME; use lance_arrow::{ARROW_EXT_NAME_KEY, iter_str_array}; use lance_bitpacking::{BitPacker, BitPacker4x}; @@ -27,18 +27,16 @@ use lance_core::deepsize::DeepSizeOf; use lance_core::error::LanceOptionExt; use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tokio::{IO_CORE_RESERVATION, get_num_compute_intensive_cpus, spawn_cpu}; -use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_core::{Error, ROW_ID, Result}; use lance_io::object_store::ObjectStore; use lance_select::RowSetOps; use object_store::path::Path; use roaring::RoaringBitmap; use smallvec::SmallVec; use std::collections::HashMap; -use std::pin::Pin; use std::str::FromStr; use std::sync::Arc; use std::sync::LazyLock; -use std::task::{Context, Poll}; use std::{fmt::Debug, sync::atomic::AtomicU64}; use tracing::instrument; @@ -1264,6 +1262,11 @@ struct WorkerOutput { tail_partition: Option, } +enum DocumentSource<'a> { + Text(&'a str), + StringList(&'a dyn Array), +} + #[derive(Debug, Clone, Copy)] struct IndexWorkerConfig { with_position: bool, @@ -1349,147 +1352,256 @@ impl IndexWorker { async fn process_batch(&mut self, batch: RecordBatch) -> Result<()> { let doc_col = batch.column(0); - let doc_iter = iter_str_array(doc_col); let row_id_col = batch[ROW_ID].as_primitive::(); - let docs = doc_iter - .zip(row_id_col.values().iter()) - .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); + match doc_col.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + let docs = iter_str_array(doc_col.as_ref()) + .zip(row_id_col.values().iter()) + .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); + + for (doc, row_id) in docs { + self.process_document(row_id, DocumentSource::Text(doc), false) + .await?; + } + } + DataType::List(_) => { + self.process_string_list_batch::(doc_col, row_id_col) + .await?; + } + DataType::LargeList(_) => { + self.process_string_list_batch::(doc_col, row_id_col) + .await?; + } + data_type => { + return Err(Error::index(format!( + "expect data type String, LargeString, List(String), or LargeList(String) but got {}", + data_type + ))); + } + } + + Ok(()) + } + async fn process_string_list_batch( + &mut self, + doc_col: &Arc, + row_id_col: &arrow_array::PrimitiveArray, + ) -> Result<()> { + let docs = doc_col.as_list::(); + match docs.value_type() { + datatypes::DataType::Utf8 | datatypes::DataType::LargeUtf8 => {} + data_type => { + return Err(Error::index(format!( + "expect list item data type String or LargeString but got {}", + data_type + ))); + } + } + + for (doc, row_id) in docs.iter().zip(row_id_col.values().iter()) { + let Some(doc) = doc else { + continue; + }; + + self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), true) + .await?; + } + + Ok(()) + } + + fn checked_token_position(row_id: u64, token_position: usize) -> Result { + u32::try_from(token_position).map_err(|_| { + Error::invalid_input(format!( + "token position overflow for row_id={row_id}: token_position={token_position}" + )) + }) + } + + fn materialize_string_list(elements: &dyn Array) -> String { + let mut doc = String::new(); + for element in iter_str_array(elements).flatten() { + if !doc.is_empty() { + doc.push(' '); + } + doc.push_str(element); + } + doc + } + + async fn process_document( + &mut self, + row_id: u64, + document: DocumentSource<'_>, + skip_empty_document: bool, + ) -> Result<()> { let with_position = self.has_position(); - for (doc, row_id) in docs { - let builder_was_empty = self.builder.docs.is_empty(); - let old_temporary_memory_size = self.temporary_memory_size(); - let old_token_memory_size = self.builder.tokens.memory_size() as u64; - let doc_id = self.builder.docs.len() as u32; - let mut token_num: u32 = 0; - let mut posting_memory_delta = 0i64; - if with_position { + let builder_was_empty = self.builder.docs.is_empty(); + let old_temporary_memory_size = self.temporary_memory_size(); + let old_token_memory_size = self.builder.tokens.memory_size() as u64; + let doc_id = self.builder.docs.len() as u32; + let mut token_num: u32 = 0; + let mut doc_length_bytes = 0usize; + let mut posting_memory_delta = 0i64; + if with_position { + { if self.token_ids.capacity() < self.last_token_count { self.token_ids .reserve(self.last_token_count - self.token_ids.capacity()); } 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 mut token_stream = self.tokenizer.token_stream_for_doc(doc); - while token_stream.advance() { - let token = token_stream.token(); - 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( - true, - posting_tail_codec, - ), - ); - 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 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( + true, + posting_tail_codec, + ), + ); + 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; } - 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, token.position as u32)? { - token_ids.push(token_id); + Ok(()) + }; + + match document { + DocumentSource::Text(doc) => { + process_text(doc)?; + } + DocumentSource::StringList(elements) => { + let doc = Self::materialize_string_list(elements); + process_text(&doc)?; } - 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 { + } + } 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 mut token_stream = self.tokenizer.token_stream_for_doc(doc); - while token_stream.advance() { - let token_id = self.builder.tokens.get_or_add(&token_stream.token().text); - self.token_ids.push(token_id); - token_num += 1; - } - } - self.adjust_tracked_memory_size( - old_token_memory_size, - self.builder.tokens.memory_size() as u64, - ); + let tokenizer = &mut self.tokenizer; + let builder = &mut self.builder; + let token_ids = &mut self.token_ids; + let mut process_text = |text: &str| { + doc_length_bytes += text.len(); + let mut token_stream = tokenizer.token_stream_for_doc(text); + while token_stream.advance() { + let token_id = builder.tokens.get_or_add(&token_stream.token().text); + token_ids.push(token_id); + token_num += 1; + } + }; - if !with_position { - let old_posting_lists_overhead_size = self.posting_lists_overhead_size(); - self.builder - .posting_lists - .resize_with(self.builder.tokens.len(), || { - PostingListBuilder::new_with_posting_tail_codec( - false, - self.builder.posting_tail_codec, - ) - }); - let new_posting_lists_overhead_size = self.posting_lists_overhead_size(); - Self::adjust_tracked_value( - &mut self.memory_size, - old_posting_lists_overhead_size, - new_posting_lists_overhead_size, - ); + match document { + DocumentSource::Text(doc) => process_text(doc), + DocumentSource::StringList(elements) => { + let doc = Self::materialize_string_list(elements); + process_text(&doc); + } + } } + } + self.adjust_tracked_memory_size( + old_token_memory_size, + self.builder.tokens.memory_size() as u64, + ); - let old_doc_memory_size = self.builder.docs.memory_size() as u64; - let appended_doc_id = self.builder.docs.append(row_id, token_num); - debug_assert_eq!(appended_doc_id, doc_id); + if skip_empty_document && token_num == 0 { + self.last_token_count = 0; + self.trim_temporary_buffers(); self.adjust_tracked_memory_size( - old_doc_memory_size, - self.builder.docs.memory_size() as u64, + old_temporary_memory_size, + self.temporary_memory_size(), ); - self.total_doc_length += doc.len(); + return Ok(()); + } - if with_position { - for &token_id in &self.token_ids { - let (old_posting_memory_size, new_posting_memory_size) = { - let posting_list = &mut self.builder.posting_lists[token_id as usize]; - let old_posting_memory_size = posting_list.size(); - posting_list.finish_open_doc(doc_id)?; - let new_posting_memory_size = posting_list.size(); - (old_posting_memory_size, new_posting_memory_size) - }; - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; - } - Self::apply_delta(&mut self.memory_size, posting_memory_delta); - } else if token_num > 0 { - self.token_ids.sort_unstable(); - let mut iter = self.token_ids.iter(); - let mut current = *iter.next().unwrap(); - let mut count = 1u32; - for &token_id in iter { - if token_id == current { - count += 1; - continue; - } + if !with_position { + let old_posting_lists_overhead_size = self.posting_lists_overhead_size(); + self.builder + .posting_lists + .resize_with(self.builder.tokens.len(), || { + PostingListBuilder::new_with_posting_tail_codec( + false, + self.builder.posting_tail_codec, + ) + }); + let new_posting_lists_overhead_size = self.posting_lists_overhead_size(); + Self::adjust_tracked_value( + &mut self.memory_size, + old_posting_lists_overhead_size, + new_posting_lists_overhead_size, + ); + } - let (old_posting_memory_size, new_posting_memory_size) = { - let posting_list = &mut self.builder.posting_lists[current as usize]; - let old_posting_memory_size = posting_list.size(); - posting_list.add(doc_id, PositionRecorder::Count(count)); - let new_posting_memory_size = posting_list.size(); - (old_posting_memory_size, new_posting_memory_size) - }; - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; + let old_doc_memory_size = self.builder.docs.memory_size() as u64; + let appended_doc_id = self.builder.docs.append(row_id, token_num); + debug_assert_eq!(appended_doc_id, doc_id); + self.adjust_tracked_memory_size( + old_doc_memory_size, + self.builder.docs.memory_size() as u64, + ); + self.total_doc_length += doc_length_bytes; - current = token_id; - count = 1; + if with_position { + for &token_id in &self.token_ids { + let (old_posting_memory_size, new_posting_memory_size) = { + let posting_list = &mut self.builder.posting_lists[token_id as usize]; + let old_posting_memory_size = posting_list.size(); + posting_list.finish_open_doc(doc_id)?; + let new_posting_memory_size = posting_list.size(); + (old_posting_memory_size, new_posting_memory_size) + }; + posting_memory_delta += + new_posting_memory_size as i64 - old_posting_memory_size as i64; + } + Self::apply_delta(&mut self.memory_size, posting_memory_delta); + } else if token_num > 0 { + self.token_ids.sort_unstable(); + let mut iter = self.token_ids.iter(); + let mut current = *iter.next().unwrap(); + let mut count = 1u32; + for &token_id in iter { + if token_id == current { + count += 1; + continue; } + let (old_posting_memory_size, new_posting_memory_size) = { let posting_list = &mut self.builder.posting_lists[current as usize]; let old_posting_memory_size = posting_list.size(); @@ -1499,27 +1611,35 @@ impl IndexWorker { }; posting_memory_delta += new_posting_memory_size as i64 - old_posting_memory_size as i64; - Self::apply_delta(&mut self.memory_size, posting_memory_delta); - } - self.last_token_count = self.token_ids.len(); - self.trim_temporary_buffers(); - self.adjust_tracked_memory_size( - old_temporary_memory_size, - self.temporary_memory_size(), - ); - if self.builder.docs.len() == 1 && self.memory_size > self.worker_memory_limit_bytes { - return Err(Error::invalid_input(format!( - "single document row_id={} exceeds worker memory limit: {} > {} bytes", - row_id, self.memory_size, self.worker_memory_limit_bytes - ))); + current = token_id; + count = 1; } + let (old_posting_memory_size, new_posting_memory_size) = { + let posting_list = &mut self.builder.posting_lists[current as usize]; + let old_posting_memory_size = posting_list.size(); + posting_list.add(doc_id, PositionRecorder::Count(count)); + let new_posting_memory_size = posting_list.size(); + (old_posting_memory_size, new_posting_memory_size) + }; + posting_memory_delta += new_posting_memory_size as i64 - old_posting_memory_size as i64; + Self::apply_delta(&mut self.memory_size, posting_memory_delta); + } + self.last_token_count = self.token_ids.len(); + self.trim_temporary_buffers(); + self.adjust_tracked_memory_size(old_temporary_memory_size, self.temporary_memory_size()); - if self.builder.docs.len() as u32 == u32::MAX - || (!builder_was_empty && self.memory_size >= self.worker_memory_limit_bytes) - { - self.flush().await?; - } + if self.builder.docs.len() == 1 && self.memory_size > self.worker_memory_limit_bytes { + return Err(Error::invalid_input(format!( + "single document row_id={} exceeds worker memory limit: {} > {} bytes", + row_id, self.memory_size, self.worker_memory_limit_bytes + ))); + } + + if self.builder.docs.len() as u32 == u32::MAX + || (!builder_was_empty && self.memory_size >= self.worker_memory_limit_bytes) + { + self.flush().await?; } Ok(()) @@ -1772,129 +1892,6 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( Arc::new(arrow_schema::Schema::new_with_metadata(fields, metadata)) } -/// Flatten the string list stream into a string stream -pub struct FlattenStream { - /// Inner record batch stream with 2 columns: - /// 1. doc_col: List(Utf8) or List(LargeUtf8) - /// 2. row_id_col: UInt64 - inner: SendableRecordBatchStream, - field_type: DataType, - data_type: DataType, -} - -impl FlattenStream { - pub fn new(input: SendableRecordBatchStream) -> Self { - let schema = input.schema(); - let field = schema.field(0); - let data_type = match field.data_type() { - DataType::List(f) if matches!(f.data_type(), DataType::Utf8) => DataType::Utf8, - DataType::List(f) if matches!(f.data_type(), DataType::LargeUtf8) => { - DataType::LargeUtf8 - } - DataType::LargeList(f) if matches!(f.data_type(), DataType::Utf8) => DataType::Utf8, - DataType::LargeList(f) if matches!(f.data_type(), DataType::LargeUtf8) => { - DataType::LargeUtf8 - } - _ => panic!( - "expect data type List(Utf8) or List(LargeUtf8) but got {:?}", - field.data_type() - ), - }; - Self { - inner: input, - field_type: field.data_type().clone(), - data_type, - } - } -} - -impl Stream for FlattenStream { - type Item = datafusion_common::Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match Pin::new(&mut self.inner).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => { - let doc_col = batch.column(0); - let batch = match self.field_type { - DataType::List(_) => flatten_string_list::(&batch, doc_col).map_err(|e| { - datafusion_common::error::DataFusionError::Execution(format!( - "flatten string list error: {}", - e - )) - }), - DataType::LargeList(_) => { - flatten_string_list::(&batch, doc_col).map_err(|e| { - datafusion_common::error::DataFusionError::Execution(format!( - "flatten string list error: {}", - e - )) - }) - } - _ => unreachable!( - "expect data type List or LargeList but got {:?}", - self.field_type - ), - }; - Poll::Ready(Some(batch)) - } - Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), - Poll::Ready(None) => Poll::Ready(None), - Poll::Pending => Poll::Pending, - } - } -} - -impl RecordBatchStream for FlattenStream { - fn schema(&self) -> SchemaRef { - let schema = Schema::new(vec![ - Field::new( - self.inner.schema().field(0).name(), - self.data_type.clone(), - true, - ), - ROW_ID_FIELD.clone(), - ]); - - Arc::new(schema) - } -} - -fn flatten_string_list( - batch: &RecordBatch, - doc_col: &Arc, -) -> Result { - let docs = doc_col.as_list::(); - let row_ids = batch[ROW_ID].as_primitive::(); - - let row_ids = row_ids - .values() - .iter() - .zip(docs.iter()) - .flat_map(|(row_id, doc)| std::iter::repeat_n(*row_id, doc.map(|d| d.len()).unwrap_or(0))); - - let row_ids = Arc::new(UInt64Array::from_iter_values(row_ids)); - let docs = match docs.value_type() { - datatypes::DataType::Utf8 | datatypes::DataType::LargeUtf8 => docs.values().clone(), - _ => { - return Err(Error::index(format!( - "expect data type String or LargeString but got {}", - docs.value_type() - ))); - } - }; - - let schema = Schema::new(vec![ - Field::new( - batch.schema().field(0).name(), - docs.data_type().clone(), - true, - ), - ROW_ID_FIELD.clone(), - ]); - let batch = RecordBatch::try_new(Arc::new(schema), vec![docs, row_ids])?; - Ok(batch) -} - pub(crate) fn token_file_path(partition_id: u64) -> String { format!("part_{}_{}", partition_id, TOKENS_FILE) } @@ -2168,7 +2165,7 @@ pub fn document_input( DataType::List(field) | DataType::LargeList(field) if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) => { - Ok(Box::pin(FlattenStream::new(input))) + Ok(input) } DataType::LargeBinary => match field.metadata().get(ARROW_EXT_NAME_KEY) { Some(name) if name.as_str() == JSON_EXT_NAME => { diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index d5a513c7112..74bc77fbbcf 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -5052,12 +5052,11 @@ impl DocSet { /// Resolve a `row_id` to every `doc_id` it owns. /// - /// A scalar column maps each row to a single document, but a - /// `list` column indexes every element as its own document, so a - /// single `row_id` can own several `doc_id`s sharing that key in `inv`. + /// Modern indexes map each row to a single document. Older list indexes + /// may have indexed each list element as its own document, so a single + /// `row_id` can still own several `doc_id`s sharing that key in `inv`. /// The prefilter path (`flat_search`) walks an allow-list of row_ids and - /// must evaluate *all* of a row's documents; resolving to one `doc_id` - /// silently drops matches at non-last list positions (lancedb#3352). + /// must evaluate all legacy documents for that row. pub fn doc_ids(&self, row_id: u64) -> impl Iterator + '_ { if self.inv.is_empty() { // in legacy format, the row id is doc id (one document per row) @@ -5330,6 +5329,12 @@ pub fn flat_full_text_search( match batches[0][doc_col].data_type() { DataType::Utf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), DataType::LargeUtf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), + DataType::List(_) => { + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + } + DataType::LargeList(_) => { + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + } data_type => Err(Error::invalid_input(format!( "unsupported data type {} for inverted index", data_type @@ -5365,6 +5370,46 @@ fn do_flat_full_text_search( Ok(results) } +fn do_flat_full_text_search_list( + batches: &[&RecordBatch], + doc_col: &str, + query: &str, + tokenizer: Option>, +) -> Result> { + let mut results = Vec::new(); + let mut tokenizer = + tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap()); + let query_tokens = collect_query_tokens(query, &mut tokenizer); + + for batch in batches { + let row_id_array = batch[ROW_ID].as_primitive::(); + let doc_array = batch[doc_col].as_list::(); + match doc_array.value_type() { + DataType::Utf8 | DataType::LargeUtf8 => {} + data_type => { + return Err(Error::invalid_input(format!( + "unsupported list item data type {} for inverted index", + data_type + ))); + } + } + for i in 0..row_id_array.len() { + if doc_array.is_null(i) { + continue; + } + let elements = doc_array.value(i); + if iter_str_array(elements.as_ref()) + .flatten() + .any(|element| has_query_token(element, &mut tokenizer, &query_tokens)) + { + results.push(row_id_array.value(i)); + } + } + } + + Ok(results) +} + const FLAT_ROW_ID_COL_IDX: usize = 0; const FLAT_ALL_TOKENS_COL_IDX: usize = 1; const FLAT_QUERY_TOKEN_COUNTS_COL_IDX: usize = 2; @@ -5422,6 +5467,8 @@ async fn tokenize_and_count( // thread is invisible to the caller's poll timer otherwise). let start = std::time::Instant::now(); let batch = batch?; + let row_id_array = batch[ROW_ID].as_primitive::(); + let mut row_ids = UInt64Builder::with_capacity(batch.num_rows()); let mut all_token_counts = UInt64Builder::with_capacity(batch.num_rows()); let mut query_token_counts = FixedSizeListBuilder::with_capacity( UInt64Builder::with_capacity(batch.num_rows() * query_tokens.len()), @@ -5429,20 +5476,7 @@ async fn tokenize_and_count( batch.num_rows(), ); let mut temp_query_token_counts = Vec::with_capacity(query_tokens.len()); - let doc_iter = iter_str_array(batch.column(doc_col_idx)); - for doc in doc_iter { - let Some(doc) = doc else { - all_token_counts.append_value(0); - query_token_counts - .values() - .append_value_n(0, query_tokens.len()); - query_token_counts.append(true); - continue; - }; - - temp_query_token_counts.clear(); - temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens.len())); - + let mut count_text = |doc: &str, temp_query_token_counts: &mut Vec| -> u64 { let mut stream = tokenizer.token_stream_for_doc(doc); let mut all_tokens = 0; while let Some(token) = stream.next() { @@ -5451,20 +5485,67 @@ async fn tokenize_and_count( temp_query_token_counts[token_index] += 1; } } - all_token_counts.append_value(all_tokens); - for count in temp_query_token_counts.iter().copied() { - query_token_counts.values().append_value(count); + all_tokens + }; + let mut append_counts = + |row_id: u64, all_tokens: u64, temp_query_token_counts: &[u64]| { + row_ids.append_value(row_id); + all_token_counts.append_value(all_tokens); + for count in temp_query_token_counts.iter().copied() { + query_token_counts.values().append_value(count); + } + query_token_counts.append(true); + }; + match batch.column(doc_col_idx).data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + let doc_iter = iter_str_array(batch.column(doc_col_idx)); + for (doc, row_id) in doc_iter.zip(row_id_array.values().iter()) { + temp_query_token_counts.clear(); + temp_query_token_counts + .extend(std::iter::repeat_n(0, query_tokens.len())); + + let Some(doc) = doc else { + append_counts(*row_id, 0, &temp_query_token_counts); + continue; + }; + + let all_tokens = count_text(doc, &mut temp_query_token_counts); + append_counts(*row_id, all_tokens, &temp_query_token_counts); + } + } + DataType::List(_) => { + tokenize_and_count_list::( + batch.column(doc_col_idx), + row_id_array, + &mut count_text, + &mut append_counts, + &mut temp_query_token_counts, + query_tokens.len(), + )?; + } + DataType::LargeList(_) => { + tokenize_and_count_list::( + batch.column(doc_col_idx), + row_id_array, + &mut count_text, + &mut append_counts, + &mut temp_query_token_counts, + query_tokens.len(), + )?; + } + data_type => { + return DataFusionResult::Err(datafusion_common::DataFusionError::Execution( + format!("unsupported data type {} for flat full text search", data_type), + )); } - query_token_counts.append(true); } - let row_ids = batch[ROW_ID].clone(); + let row_ids = row_ids.finish(); let all_token_counts = all_token_counts.finish(); let query_token_counts = query_token_counts.finish(); let result_batch = RecordBatch::try_new( - output_schema, vec![ - row_ids, + Arc::new(row_ids) as ArrayRef, Arc::new(all_token_counts) as ArrayRef, Arc::new(query_token_counts) as ArrayRef, ], @@ -5490,6 +5571,47 @@ async fn tokenize_and_count( )?) } +fn tokenize_and_count_list( + doc_col: &ArrayRef, + row_id_array: &arrow_array::PrimitiveArray, + count_text: &mut impl FnMut(&str, &mut Vec) -> u64, + append_counts: &mut impl FnMut(u64, u64, &[u64]), + temp_query_token_counts: &mut Vec, + query_tokens_len: usize, +) -> DataFusionResult<()> { + let doc_array = doc_col.as_list::(); + match doc_array.value_type() { + DataType::Utf8 | DataType::LargeUtf8 => {} + data_type => { + return Err(datafusion_common::DataFusionError::Execution(format!( + "unsupported list item data type {} for flat full text search", + data_type + ))); + } + } + + for i in 0..row_id_array.len() { + if doc_array.is_null(i) { + continue; + } + + temp_query_token_counts.clear(); + temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens_len)); + + let elements = doc_array.value(i); + let mut all_tokens = 0; + for element in iter_str_array(elements.as_ref()).flatten() { + all_tokens += count_text(element, temp_query_token_counts); + } + + if all_tokens > 0 { + append_counts(row_id_array.value(i), all_tokens, temp_query_token_counts); + } + } + + Ok(()) +} + /// Initialize the BM25 scorer /// /// In order to calculate BM25 scores we need to know token counts for the entire corpus. We extract these from the @@ -5530,7 +5652,9 @@ fn initialize_scorer( for _ in 0..counted_input.num_rows() { for token_count in all_token_counts.iter_mut() { - *token_count += input_token_counters.next().unwrap_or_default(); + if input_token_counters.next().unwrap_or_default() > 0 { + *token_count += 1; + } } } @@ -5742,7 +5866,10 @@ mod tests { }; use crate::scalar::inverted::query::{FtsSearchParams, Operator}; use crate::scalar::lance_format::LanceIndexStore; - use arrow::array::{AsArray, Int32Builder, LargeBinaryBuilder, ListBuilder, UInt32Builder}; + use arrow::array::{ + AsArray, GenericListBuilder, GenericStringBuilder, Int32Builder, LargeBinaryBuilder, + ListBuilder, UInt32Builder, + }; use arrow::datatypes::{Float32Type, UInt32Type}; use arrow_array::{ArrayRef, Float32Array, RecordBatch, StringArray, UInt32Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; @@ -8815,6 +8942,60 @@ mod tests { ); } + #[tokio::test] + async fn flat_bm25_search_treats_string_lists_as_row_documents() { + let mut docs_builder = + GenericListBuilder::::new(GenericStringBuilder::::new()); + docs_builder.values().append_value("alpha"); + docs_builder.values().append_value("alpha beta"); + docs_builder.append(true); + docs_builder.values().append_value("beta"); + docs_builder.append(true); + docs_builder.append(true); + docs_builder.values().append_null(); + docs_builder.append(true); + docs_builder.append(false); + + let docs = Arc::new(docs_builder.finish()) as ArrayRef; + let schema = Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("text", docs.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(vec![0u64, 1, 2, 3, 4])) as ArrayRef, + docs, + ], + ) + .unwrap(); + + let input: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + stream::iter(vec![Ok(batch)]), + )); + let tokenizer: Box = Box::new(TextTokenizer::new( + TextAnalyzer::builder(SimpleTokenizer::default()).build(), + )); + + let result_stream = flat_bm25_search_stream_with_metrics( + input, + "text".to_string(), + "alpha".to_string(), + tokenizer, + None, + 100, + None, + ) + .await + .unwrap(); + let batches: Vec<_> = result_stream.try_collect().await.unwrap(); + let scored = arrow::compute::concat_batches(&FTS_SCHEMA, &batches).unwrap(); + let row_ids = scored[ROW_ID].as_primitive::(); + + assert_eq!(row_ids.values(), &[0]); + } + /// An [`IndexReader`] wrapper that hides the posting-group-offsets schema /// metadata key, so a [`PostingListReader`] opened on it takes the /// pre-grouping per-token fallback path (issue #7040). diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 961b381e452..de79a321a3e 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1803,6 +1803,169 @@ async fn test_fts_index_with_large_string() { test_fts_index::(true).await; } +#[tokio::test] +async fn test_fts_list_index_uses_row_level_documents() { + let tempdir = TempStrDir::default(); + let uri = tempdir.to_owned(); + drop(tempdir); + + let mut list_col = GenericListBuilder::::new(GenericStringBuilder::::new()); + list_col.values().append_value("lance"); + list_col.values().append_value("lance database"); + list_col.append(true); + list_col.values().append_value("database"); + list_col.append(true); + list_col.append(true); + list_col.values().append_null(); + list_col.append(true); + list_col.append(false); + + let docs = Arc::new(list_col.finish()) as ArrayRef; + let ids = Arc::new(UInt64Array::from_iter_values(0..docs.len() as u64)) as ArrayRef; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("doc", docs.data_type().clone(), true), + ArrowField::new("id", DataType::UInt64, false), + ])), + vec![docs, ids], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + + dataset + .create_index( + &["doc"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("lance".to_owned()).limit(Some(10))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[0]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("database".to_owned()).limit(Some(10))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut ids = result["id"] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1], "{:?}", result); +} + +#[tokio::test] +async fn test_fts_list_phrase_query_can_cross_elements() { + assert_fts_list_phrase_query_can_cross_elements::().await; +} + +#[tokio::test] +async fn test_fts_large_list_phrase_query_can_cross_elements() { + assert_fts_list_phrase_query_can_cross_elements::().await; +} + +async fn assert_fts_list_phrase_query_can_cross_elements() { + let tempdir = TempStrDir::default(); + let uri = tempdir.to_owned(); + drop(tempdir); + + let mut list_col = GenericListBuilder::::new(GenericStringBuilder::::new()); + let rows: &[&[&str]] = &[ + &["alpha", "beta"], + &["want the", "apple"], + &["want", "apple"], + ]; + for values in rows.iter().copied() { + for value in values { + list_col.values().append_value(value); + } + list_col.append(true); + } + + let docs = Arc::new(list_col.finish()) as ArrayRef; + let ids = Arc::new(UInt64Array::from(vec![0u64, 1, 2])) as ArrayRef; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("doc", docs.data_type().clone(), true), + ArrowField::new("id", DataType::UInt64, false), + ])), + vec![docs, ids], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + + let params = InvertedIndexParams::default() + .with_position(true) + .remove_stop_words(true); + dataset + .create_index(&["doc"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("alpha beta".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[0]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("want the apple".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[1]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("want apple".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[2]); +} + #[tokio::test] async fn test_fts_accented_chars() { let ds = create_fts_dataset::(false, false, InvertedIndexParams::default()).await;