From 8f5d784e5acca2ecf3d1df092dfd45a6ddba41a8 Mon Sep 17 00:00:00 2001 From: Du Bin Date: Thu, 20 Aug 2026 09:04:36 +0000 Subject: [PATCH] feat(index): support Utf8View in FTS indices --- rust/lance-arrow/src/lib.rs | 18 ++++++++-- rust/lance-index/src/scalar/inverted.rs | 36 ++++++++++++++++--- .../src/scalar/inverted/builder.rs | 19 ++++++---- .../src/scalar/inverted/index/flat_search.rs | 23 ++++++------ .../inverted/index/tests/flat_search.rs | 19 ++++++++++ .../inverted/tokenizer/document_tokenizer.rs | 29 +++++++++++++-- 6 files changed, 116 insertions(+), 28 deletions(-) diff --git a/rust/lance-arrow/src/lib.rs b/rust/lance-arrow/src/lib.rs index a08ad4c67d5..f3bb8a061d1 100644 --- a/rust/lance-arrow/src/lib.rs +++ b/rust/lance-arrow/src/lib.rs @@ -471,7 +471,11 @@ pub fn iter_str_array(arr: &dyn Array) -> Box> match arr.data_type() { DataType::Utf8 => Box::new(arr.as_string::().iter()), DataType::LargeUtf8 => Box::new(arr.as_string::().iter()), - _ => panic!("Expecting Utf8 or LargeUtf8, found {:?}", arr.data_type()), + DataType::Utf8View => Box::new(arr.as_string_view().iter()), + _ => panic!( + "Expecting Utf8, LargeUtf8, or Utf8View, found {:?}", + arr.data_type() + ), } } @@ -1585,9 +1589,19 @@ impl BufferExt for arrow_buffer::Buffer { mod tests { use super::*; use arrow_array::{Float32Array, Int32Array, NullArray, StructArray}; - use arrow_array::{ListArray, StringArray, new_empty_array, new_null_array}; + use arrow_array::{ListArray, StringArray, StringViewArray, new_empty_array, new_null_array}; use arrow_buffer::OffsetBuffer; + #[test] + fn test_iter_str_array_utf8_view() { + let array = StringViewArray::from(vec![Some("alpha"), None, Some("beta")]); + + assert_eq!( + iter_str_array(&array).collect::>(), + vec![Some("alpha"), None, Some("beta")] + ); + } + #[test] fn test_convert_to_floating_point_preserves_inner_nulls() { // A FixedSizeList with a null inner element must convert to a diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index b07ee2b5e73..0c52b27f042 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -216,12 +216,20 @@ impl BasicTrainer for InvertedIndexPlugin { field: &Field, ) -> Result> { match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 | DataType::LargeBinary => (), - DataType::List(f) if matches!(f.data_type(), DataType::Utf8 | DataType::LargeUtf8) => (), - DataType::LargeList(f) if matches!(f.data_type(), DataType::Utf8 | DataType::LargeUtf8) => (), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View | DataType::LargeBinary => (), + DataType::List(f) + if matches!( + f.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) => (), + DataType::LargeList(f) + if matches!( + f.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) => (), _ => return Err(Error::invalid_input_source(format!( - "A inverted index can only be created on a Utf8 or LargeUtf8 field/list or LargeBinary field. Column has type {:?}", + "An inverted index can only be created on a Utf8, LargeUtf8, or Utf8View field/list or LargeBinary field. Column has type {:?}", field.data_type() ) .into())) @@ -350,6 +358,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { mod tests { use super::*; use crate::scalar::{BuiltinIndexType, ScalarIndexParams}; + use rstest::rstest; #[test] fn test_plugin_version_tracks_v3_capability_gate() { @@ -357,6 +366,25 @@ mod tests { assert_eq!(plugin.version(), INVERTED_INDEX_VERSION_V3); } + #[rstest] + #[case::utf8_view(DataType::Utf8View)] + #[case::list_utf8_view(DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8View, + true, + ))))] + #[case::large_list_utf8_view(DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Utf8View, + true, + ))))] + fn test_new_training_request_supports_utf8_view(#[case] data_type: DataType) { + let plugin = InvertedIndexPlugin; + let field = Field::new("text", data_type, true); + + plugin.new_training_request("{}", &field).unwrap(); + } + #[test] fn test_details_json_includes_document_granularity() { let details = pbold::InvertedIndexDetails { diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index f50bd827e20..92e9cf6a8dc 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1458,7 +1458,7 @@ impl IndexWorker { } data_type => { return Err(Error::index(format!( - "expect data type String, LargeString, List(String), or LargeList(String) but got {}", + "expect String, LargeString, StringView, List(String), or LargeList(String) but got {}", data_type ))); } @@ -1474,10 +1474,12 @@ impl IndexWorker { ) -> Result<()> { let docs = doc_col.as_list::(); match docs.value_type() { - datatypes::DataType::Utf8 | datatypes::DataType::LargeUtf8 => {} + datatypes::DataType::Utf8 + | datatypes::DataType::LargeUtf8 + | datatypes::DataType::Utf8View => {} data_type => { return Err(Error::index(format!( - "expect list item data type String or LargeString but got {}", + "expect list item data type String, LargeString, or StringView but got {}", data_type ))); } @@ -2343,8 +2345,8 @@ async fn merge_metadata_files( /// Convert input stream into a stream of documents. /// /// The input stream must be one of: -/// 1. Document in Utf8 or LargeUtf8 format. -/// 2. Document in List(Utf8) or List(LargeUtf8) format. +/// 1. Document in Utf8, LargeUtf8, or Utf8View format. +/// 2. Document in List(Utf8), List(LargeUtf8), or List(Utf8View) format. /// 3. Json document in LargeBinary format. pub fn document_input( input: SendableRecordBatchStream, @@ -2355,7 +2357,10 @@ pub fn document_input( match field.data_type() { DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(input), DataType::List(field) | DataType::LargeList(field) - if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) => + if matches!( + field.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) => { Ok(input) } @@ -2369,7 +2374,7 @@ pub fn document_input( }, _ => Err(Error::invalid_input_source( format!( - "column {} has type {}, is not utf8, large utf8 type/list, or large binary", + "column {} has type {}, is not utf8, large utf8, utf8 view type/list, or large binary", column, field.data_type() ) 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..e99bac2de12 100644 --- a/rust/lance-index/src/scalar/inverted/index/flat_search.rs +++ b/rust/lance-index/src/scalar/inverted/index/flat_search.rs @@ -33,11 +33,8 @@ 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, phrase_slop) - } - DataType::LargeUtf8 => { - do_flat_full_text_search::(batches, doc_col, query, tokenizer, phrase_slop) + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + do_flat_full_text_search(batches, doc_col, query, tokenizer, phrase_slop) } DataType::List(_) => { do_flat_full_text_search_list::(batches, doc_col, query, tokenizer, phrase_slop) @@ -52,7 +49,7 @@ pub fn flat_full_text_search( } } -pub(super) fn do_flat_full_text_search( +pub(super) fn do_flat_full_text_search( batches: &[&RecordBatch], doc_col: &str, query: &str, @@ -66,11 +63,13 @@ pub(super) fn do_flat_full_text_search( for batch in batches { 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() { - let doc = doc_array.value(i); + let doc_array = &batch[doc_col]; + for (doc, row_id) in iter_str_array(doc_array).zip(row_id_array.values()) { + let Some(doc) = doc else { + continue; + }; if document_matches_flat_query(doc, &mut tokenizer, &query_tokens, phrase_slop)? { - results.push(row_id_array.value(i)); + results.push(*row_id); } } } @@ -94,7 +93,7 @@ pub(super) fn do_flat_full_text_search_list( 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 => {} + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {} data_type => { return Err(Error::invalid_input(format!( "unsupported list item data type {} for inverted index", @@ -476,7 +475,7 @@ pub(super) fn tokenize_and_count_list( ) -> DataFusionResult<()> { let doc_array = doc_col.as_list::(); match doc_array.value_type() { - DataType::Utf8 | DataType::LargeUtf8 => {} + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {} data_type => { return Err(datafusion_common::DataFusionError::Execution(format!( "unsupported list item data type {} for flat full text search", diff --git a/rust/lance-index/src/scalar/inverted/index/tests/flat_search.rs b/rust/lance-index/src/scalar/inverted/index/tests/flat_search.rs index 910678adb72..a6f13d1995d 100644 --- a/rust/lance-index/src/scalar/inverted/index/tests/flat_search.rs +++ b/rust/lance-index/src/scalar/inverted/index/tests/flat_search.rs @@ -2,6 +2,25 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use super::*; +use arrow_array::record_batch; + +#[test] +fn flat_full_text_search_supports_utf8_view() { + let batch = record_batch!( + (ROW_ID, UInt64, [0, 1, 2]), + ( + "text", + Utf8View, + [Some("alpha beta"), None, Some("beta gamma")] + ) + ) + .unwrap(); + + assert_eq!( + flat_full_text_search(&[&batch], "text", "beta", None).unwrap(), + vec![0, 2] + ); +} #[tokio::test] async fn flat_bm25_search_stream_with_metrics_records_elapsed_compute() { 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..99242bfc519 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer/document_tokenizer.rs @@ -28,9 +28,12 @@ impl TryFrom<&Field> for DocType { fn try_from(field: &Field) -> Result { match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => Ok(Self::Text), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(Self::Text), DataType::List(field) | DataType::LargeList(field) - if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) => + if matches!( + field.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) => { Ok(Self::Text) } @@ -304,11 +307,31 @@ impl TokenStream for TTStream { #[cfg(test)] mod tests { use crate::scalar::inverted::tokenizer::document_tokenizer::{ - JsonTokenizer, LanceTokenizer, flatten_json, flatten_triplet, + DocType, JsonTokenizer, LanceTokenizer, flatten_json, flatten_triplet, }; + use arrow_schema::{DataType, Field}; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer, Token}; + use rstest::rstest; use serde_json::Value; + #[rstest] + #[case::utf8_view(DataType::Utf8View)] + #[case::list_utf8_view(DataType::List(std::sync::Arc::new(Field::new( + "item", + DataType::Utf8View, + true, + ))))] + #[case::large_list_utf8_view(DataType::LargeList(std::sync::Arc::new(Field::new( + "item", + DataType::Utf8View, + true, + ))))] + fn test_doc_type_supports_utf8_view(#[case] data_type: DataType) { + let field = Field::new("text", data_type, true); + + assert!(matches!(DocType::try_from(&field), Ok(DocType::Text))); + } + #[test] fn test_json_tokenizer() { let text = r#"{