Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions rust/lance-arrow/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -471,7 +471,11 @@ pub fn iter_str_array(arr: &dyn Array) -> Box<dyn Iterator<Item = Option<&str>>
match arr.data_type() {
DataType::Utf8 => Box::new(arr.as_string::<i32>().iter()),
DataType::LargeUtf8 => Box::new(arr.as_string::<i64>().iter()),
_ => panic!("Expecting Utf8 or LargeUtf8, found {:?}", arr.data_type()),
DataType::Utf8View => Box::new(arr.as_string_view().iter()),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This leaves the public distributed post-filter path unable to query the new type. FlatMatchFilterExec::{new,new_with_document_column,new_with_segments} all set resolved_field=None; that reaches rust/lance/src/io/exec/fts.rs:2013, whose dispatch accepts only Utf8 / LargeUtf8 and returns FTS document column text is not a string for Utf8View. Since this PR now admits and builds these indexes, new_with_segments—the constructor exposed for distributed planning—fails at query time.

Please route that fallback through this shared iterator (or add an equivalent Utf8View arm) and cover the public constructor.

Reproducer

On this head I added a temporary test that passed a StringViewArray stream to FlatMatchFilterExec::build_filter_stream with resolved_field=None and collected the stream. cargo test -p lance gate_reproducer_flat_match_filter_rejects_utf8_view --locked observed the execution error above; the normal new_with_resolved_field scanner path does not.

_ => panic!(
"Expecting Utf8, LargeUtf8, or Utf8View, found {:?}",
arr.data_type()
),
}
}

Expand Down Expand Up @@ -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<_>>(),
vec![Some("alpha"), None, Some("beta")]
);
}

#[test]
fn test_convert_to_floating_point_preserves_inner_nulls() {
// A FixedSizeList<Int8> with a null inner element must convert to a
Expand Down
36 changes: 32 additions & 4 deletions rust/lance-index/src/scalar/inverted.rs
Original file line number Diff line number Diff line change
Expand Up @@ -216,12 +216,20 @@ impl BasicTrainer for InvertedIndexPlugin {
field: &Field,
) -> Result<Box<dyn TrainingRequest>> {
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()))
Expand Down Expand Up @@ -350,13 +358,33 @@ 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() {
let plugin = InvertedIndexPlugin;
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 {
Expand Down
19 changes: 12 additions & 7 deletions rust/lance-index/src/scalar/inverted/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
)));
}
Expand All @@ -1474,10 +1474,12 @@ impl IndexWorker {
) -> Result<()> {
let docs = doc_col.as_list::<Offset>();
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
)));
}
Expand Down Expand Up @@ -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,
Expand All @@ -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)
}
Expand All @@ -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()
)
Expand Down
23 changes: 11 additions & 12 deletions rust/lance-index/src/scalar/inverted/index/flat_search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,8 @@ pub fn flat_full_text_search(
};

match batches[0][doc_col].data_type() {
DataType::Utf8 => {
do_flat_full_text_search::<i32>(batches, doc_col, query, tokenizer, phrase_slop)
}
DataType::LargeUtf8 => {
do_flat_full_text_search::<i64>(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::<i32>(batches, doc_col, query, tokenizer, phrase_slop)
Expand All @@ -52,7 +49,7 @@ pub fn flat_full_text_search(
}
}

pub(super) fn do_flat_full_text_search<Offset: OffsetSizeTrait>(
pub(super) fn do_flat_full_text_search(
batches: &[&RecordBatch],
doc_col: &str,
query: &str,
Expand All @@ -66,11 +63,13 @@ pub(super) fn do_flat_full_text_search<Offset: OffsetSizeTrait>(

for batch in batches {
let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();
let doc_array = batch[doc_col].as_string::<Offset>();
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);
}
}
}
Expand All @@ -94,7 +93,7 @@ pub(super) fn do_flat_full_text_search_list<ListOffset: OffsetSizeTrait>(
let row_id_array = batch[ROW_ID].as_primitive::<UInt64Type>();
let doc_array = batch[doc_col].as_list::<ListOffset>();
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",
Expand Down Expand Up @@ -476,7 +475,7 @@ pub(super) fn tokenize_and_count_list<ListOffset: OffsetSizeTrait>(
) -> DataFusionResult<()> {
let doc_array = doc_col.as_list::<ListOffset>();
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",
Expand Down
19 changes: 19 additions & 0 deletions rust/lance-index/src/scalar/inverted/index/tests/flat_search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,12 @@ impl TryFrom<&Field> for DocType {

fn try_from(field: &Field) -> Result<Self, Self::Error> {
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)
}
Expand Down Expand Up @@ -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#"{
Expand Down
Loading