#![allow(clippy::redundant_pub_crate)]
use std::{collections::BTreeMap, sync::Arc};
use arrow_array::cast::AsArray;
use arrow_array::{
Array, ArrayRef, LargeStringArray, RecordBatch, StringArray, StringViewArray, UInt32Array,
UInt64Array,
};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use datafusion::error::DataFusionError;
use datafusion::execution::SendableRecordBatchStream;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use futures::StreamExt;
use lance_core::{
ROW_ID,
datatypes::{Field, LogicalType, Schema, format_field_path, parse_field_path},
};
use lance_index::metrics::NoOpMetricsCollector;
use lance_index::pbold::InvertedIndexDetails;
use lance_index::scalar::index_files_to_table;
use lance_index::scalar::inverted::{
DocumentGranularity, InvertedIndex, InvertedIndexParams, doc_index_storage_column,
};
use lance_index::scalar::lance_format::LanceIndexStore;
use lance_index::scalar::registry::VALUE_COLUMN_NAME;
use lance_table::format::IndexMetadata;
use prost::Message;
use roaring::RoaringBitmap;
use uuid::Uuid;
use crate::{
Dataset, Error, Result,
dataset::index::LanceIndexStoreExt,
index::{DatasetIndexExt, scalar::fetch_index_details},
};
#[derive(Debug, Clone)]
enum FtsTraversal {
Text,
Struct {
child_index: usize,
child: Box<Self>,
},
List {
child: Box<Self>,
},
}
impl FtsTraversal {
fn list_depth(&self) -> usize {
match self {
Self::Text => 0,
Self::Struct { child, .. } => child.list_depth(),
Self::List { child } => 1 + child.list_depth(),
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct ResolvedFtsField {
pub final_field_id: i32,
pub root_column: String,
pub canonical_path: String,
pub document_granularity: DocumentGranularity,
traversal: FtsTraversal,
list_depth: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct FtsDocument {
pub row_index: usize,
pub text: String,
pub doc_index: Vec<u32>,
}
impl ResolvedFtsField {
pub fn has_lists(&self) -> bool {
self.list_depth > 0
}
pub fn coordinate_rank(&self) -> usize {
if self.document_granularity.is_list_element() {
self.list_depth
} else {
0
}
}
pub fn documents_from_batch(&self, batch: &RecordBatch) -> Result<Vec<FtsDocument>> {
let column = batch.column_by_name(&self.root_column).ok_or_else(|| {
Error::invalid_input(format!(
"FTS root column '{}' is missing from the input batch",
self.root_column
))
})?;
self.documents_from_array(column, batch.num_rows())
}
fn documents_from_array(&self, column: &ArrayRef, num_rows: usize) -> Result<Vec<FtsDocument>> {
let mut documents = Vec::new();
match self.document_granularity {
DocumentGranularity::Row => {
documents.reserve(num_rows);
for row_index in 0..num_rows {
let mut text = String::new();
append_row_text(&self.traversal, column.as_ref(), row_index, &mut text)?;
documents.push(FtsDocument {
row_index,
text,
doc_index: Vec::new(),
});
}
}
DocumentGranularity::ListElement => {
for row_index in 0..num_rows {
collect_element_documents(
&self.traversal,
column.as_ref(),
row_index,
row_index,
0,
self.list_depth,
&mut Vec::with_capacity(self.list_depth),
&mut documents,
)?;
}
}
}
Ok(documents)
}
}
fn string_value(array: &dyn Array, index: usize) -> Result<Option<&str>> {
if array.is_null(index) {
return Ok(None);
}
match array.data_type() {
DataType::Utf8 => Ok(Some(
array
.as_any()
.downcast_ref::<StringArray>()
.expect("Utf8 array type")
.value(index),
)),
DataType::LargeUtf8 => Ok(Some(
array
.as_any()
.downcast_ref::<LargeStringArray>()
.expect("LargeUtf8 array type")
.value(index),
)),
DataType::Utf8View => Ok(Some(
array
.as_any()
.downcast_ref::<StringViewArray>()
.expect("Utf8View array type")
.value(index),
)),
data_type => Err(Error::internal(format!(
"FTS traversal expected a string array, got {data_type}"
))),
}
}
fn append_text(output: &mut String, text: &str) {
if !output.is_empty() {
output.push(' ');
}
output.push_str(text);
}
fn append_row_text(
traversal: &FtsTraversal,
array: &dyn Array,
index: usize,
output: &mut String,
) -> Result<()> {
match traversal {
FtsTraversal::Text => {
if let Some(text) = string_value(array, index)? {
append_text(output, text);
}
}
FtsTraversal::Struct { child_index, child } => {
if !array.is_null(index) {
let structs = array.as_struct();
append_row_text(child, structs.column(*child_index).as_ref(), index, output)?;
}
}
FtsTraversal::List { child } => match array.data_type() {
DataType::List(_) => {
let lists = array.as_list::<i32>();
if !lists.is_null(index) {
let offsets = lists.value_offsets();
let start = offsets[index] as usize;
let end = offsets[index + 1] as usize;
for element_index in start..end {
append_row_text(child, lists.values().as_ref(), element_index, output)?;
}
}
}
DataType::LargeList(_) => {
let lists = array.as_list::<i64>();
if !lists.is_null(index) {
let offsets = lists.value_offsets();
let start = offsets[index] as usize;
let end = offsets[index + 1] as usize;
for element_index in start..end {
append_row_text(child, lists.values().as_ref(), element_index, output)?;
}
}
}
data_type => {
return Err(Error::internal(format!(
"FTS traversal expected a list array, got {data_type}"
)));
}
},
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn collect_element_documents(
traversal: &FtsTraversal,
array: &dyn Array,
index: usize,
row_index: usize,
list_depth: usize,
boundary_depth: usize,
doc_index: &mut Vec<u32>,
documents: &mut Vec<FtsDocument>,
) -> Result<()> {
match traversal {
FtsTraversal::Text => Err(Error::internal(
"ListElement FTS traversal did not encounter its list boundary".to_string(),
)),
FtsTraversal::Struct { child_index, child } => {
if !array.is_null(index) {
let structs = array.as_struct();
collect_element_documents(
child,
structs.column(*child_index).as_ref(),
index,
row_index,
list_depth,
boundary_depth,
doc_index,
documents,
)?;
}
Ok(())
}
FtsTraversal::List { child } => {
let mut visit = |values: &ArrayRef, start: usize, end: usize| -> Result<()> {
for (ordinal, element_index) in (start..end).enumerate() {
let ordinal = u32::try_from(ordinal).map_err(|_| {
Error::invalid_input(format!(
"FTS element ordinal overflow for row index {row_index}"
))
})?;
doc_index.push(ordinal);
if list_depth + 1 == boundary_depth {
let mut text = String::new();
append_row_text(child, values.as_ref(), element_index, &mut text)?;
documents.push(FtsDocument {
row_index,
text,
doc_index: doc_index.clone(),
});
} else {
collect_element_documents(
child,
values.as_ref(),
element_index,
row_index,
list_depth + 1,
boundary_depth,
doc_index,
documents,
)?;
}
doc_index.pop();
}
Ok(())
};
match array.data_type() {
DataType::List(_) => {
let lists = array.as_list::<i32>();
if !lists.is_null(index) {
let offsets = lists.value_offsets();
visit(
lists.values(),
offsets[index] as usize,
offsets[index + 1] as usize,
)?;
}
}
DataType::LargeList(_) => {
let lists = array.as_list::<i64>();
if !lists.is_null(index) {
let offsets = lists.value_offsets();
visit(
lists.values(),
offsets[index] as usize,
offsets[index + 1] as usize,
)?;
}
}
data_type => {
return Err(Error::internal(format!(
"FTS traversal expected a list array, got {data_type}"
)));
}
}
Ok(())
}
}
}
fn find_child_case_insensitive<'a>(field: &'a Field, name: &str) -> Option<(usize, &'a Field)> {
field
.children
.iter()
.enumerate()
.find(|(_, child)| child.name == name)
.or_else(|| {
field
.children
.iter()
.enumerate()
.find(|(_, child)| child.name.eq_ignore_ascii_case(name))
})
}
fn build_fts_traversal(
field: &Field,
remaining_names: &[String],
canonical_names: &mut Vec<String>,
final_field_id: &mut i32,
) -> Result<FtsTraversal> {
match field.data_type() {
DataType::List(_) | DataType::LargeList(_) => {
let child = field.children.first().ok_or_else(|| {
Error::invalid_input(format!(
"FTS list field '{}' does not have an item field",
field.name
))
})?;
Ok(FtsTraversal::List {
child: Box::new(build_fts_traversal(
child,
remaining_names,
canonical_names,
final_field_id,
)?),
})
}
DataType::Struct(_) => {
let Some((name, rest)) = remaining_names.split_first() else {
return Err(Error::invalid_input(format!(
"FTS field path ends at struct '{}'; specify the final text field",
field.name
)));
};
let (child_index, child) =
find_child_case_insensitive(field, name).ok_or_else(|| {
Error::index(format!(
"FTS field '{}' does not contain child '{}'",
field.name, name
))
})?;
canonical_names.push(child.name.clone());
*final_field_id = child.id;
Ok(FtsTraversal::Struct {
child_index,
child: Box::new(build_fts_traversal(
child,
rest,
canonical_names,
final_field_id,
)?),
})
}
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View if remaining_names.is_empty() => {
Ok(FtsTraversal::Text)
}
_ if field.logical_type == LogicalType::from("json") && remaining_names.is_empty() => {
Ok(FtsTraversal::Text)
}
data_type if !remaining_names.is_empty() => Err(Error::invalid_input(format!(
"FTS field '{}' has type {data_type} and cannot contain child '{}'",
field.name, remaining_names[0]
))),
data_type => Err(Error::invalid_input(format!(
"FTS field '{}' must resolve to Utf8, LargeUtf8, Utf8View, or JSON, got {data_type}",
field.name
))),
}
}
pub(crate) fn resolve_fts_field(
schema: &Schema,
path: &str,
document_granularity: DocumentGranularity,
) -> Result<ResolvedFtsField> {
let names = parse_field_path(path)?;
let (root_name, remaining_names) = names
.split_first()
.ok_or_else(|| Error::invalid_input("FTS field path cannot be empty".to_string()))?;
let root = schema
.fields
.iter()
.find(|field| field.name == *root_name)
.or_else(|| {
schema
.fields
.iter()
.find(|field| field.name.eq_ignore_ascii_case(root_name))
})
.ok_or_else(|| {
Error::index(format!(
"FTS field path '{path}' does not exist in the dataset schema"
))
})?;
let mut canonical_names = vec![root.name.clone()];
let mut final_field_id = root.id;
let traversal = build_fts_traversal(
root,
remaining_names,
&mut canonical_names,
&mut final_field_id,
)?;
let list_depth = traversal.list_depth();
if document_granularity.is_list_element() && list_depth == 0 {
return Err(Error::invalid_input(format!(
"FTS field path '{}' has no List layer and cannot use ListElement document granularity",
format_field_path(
&canonical_names
.iter()
.map(String::as_str)
.collect::<Vec<_>>()
)
)));
}
if list_depth > 0 && root.logical_type == LogicalType::from("json") {
return Err(Error::invalid_input(
"nested List traversal is not supported for JSON FTS sources".to_string(),
));
}
let canonical_path = format_field_path(
&canonical_names
.iter()
.map(String::as_str)
.collect::<Vec<_>>(),
);
Ok(ResolvedFtsField {
final_field_id,
root_column: root.name.clone(),
canonical_path,
document_granularity,
traversal,
list_depth,
})
}
fn find_public_path_by_id(field: &Field, field_id: i32, path: &mut Vec<String>) -> bool {
if field.id == field_id {
return true;
}
match field.data_type() {
DataType::List(_) | DataType::LargeList(_) => field
.children
.first()
.is_some_and(|child| find_physical_path_by_id(child, field_id, path)),
DataType::Struct(_) => field.children.iter().any(|child| {
path.push(child.name.clone());
let found = find_public_path_by_id(child, field_id, path);
if !found {
path.pop();
}
found
}),
_ => false,
}
}
fn find_physical_path_by_id(field: &Field, field_id: i32, path: &mut Vec<String>) -> bool {
if field.id == field_id {
return false;
}
match field.data_type() {
DataType::List(_) | DataType::LargeList(_) => field
.children
.first()
.is_some_and(|child| find_physical_path_by_id(child, field_id, path)),
DataType::Struct(_) => field.children.iter().any(|child| {
path.push(child.name.clone());
let found = find_public_path_by_id(child, field_id, path);
if !found {
path.pop();
}
found
}),
_ => false,
}
}
pub(crate) fn resolve_fts_field_by_id(
schema: &Schema,
field_id: i32,
document_granularity: DocumentGranularity,
) -> Result<ResolvedFtsField> {
for root in &schema.fields {
let mut path = vec![root.name.clone()];
if find_public_path_by_id(root, field_id, &mut path) {
let path_refs = path.iter().map(String::as_str).collect::<Vec<_>>();
return resolve_fts_field(schema, &format_field_path(&path_refs), document_granularity);
}
}
Err(Error::invalid_input(format!(
"FTS index refers to missing or Arrow-internal field id {field_id}"
)))
}
pub(crate) async fn indexed_fts_document_granularities(
dataset: &Dataset,
column: &str,
) -> Result<Vec<(String, DocumentGranularity)>> {
let resolved = resolve_fts_field(dataset.schema(), column, DocumentGranularity::Row)?;
let mut by_name = BTreeMap::new();
for index in dataset.load_indices().await?.iter() {
if index.keyed_field() != Some(resolved.final_field_id) {
continue;
}
let details_any = fetch_index_details(dataset, &resolved.canonical_path, index).await?;
if !details_any.type_url.ends_with("InvertedIndexDetails") {
continue;
}
let details =
InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| {
Error::corrupt_file(
dataset.indices_dir().join(index.uuid.to_string()),
format!(
"failed to decode InvertedIndexDetails for FTS index '{}': {error}",
index.name
),
)
})?;
let document_granularity = DocumentGranularity::try_from(details.document_granularity)?;
if let Some(existing) = by_name.insert(index.name.clone(), document_granularity)
&& existing != document_granularity
{
return Err(Error::corrupt_file(
dataset.indices_dir().join(index.uuid.to_string()),
format!(
"FTS index '{}' has inconsistent document granularity across segments: \
{existing:?} and {document_granularity:?}",
index.name
),
));
}
}
Ok(by_name.into_iter().collect())
}
pub(crate) async fn resolve_query_document_granularity(
dataset: &Dataset,
column: &str,
requested: Option<DocumentGranularity>,
) -> Result<DocumentGranularity> {
let indices = indexed_fts_document_granularities(dataset, column).await?;
let mut available = indices
.iter()
.map(|(_, document_granularity)| *document_granularity)
.collect::<Vec<_>>();
available.sort_by_key(|document_granularity| match document_granularity {
DocumentGranularity::Row => 0,
DocumentGranularity::ListElement => 1,
});
available.dedup();
let resolved = match requested {
Some(requested) if available.is_empty() || available.contains(&requested) => requested,
Some(requested) => {
let indexed = indices
.iter()
.map(|(name, document_granularity)| format!("'{name}' ({document_granularity:?})"))
.collect::<Vec<_>>()
.join(", ");
return Err(Error::invalid_input(format!(
"FTS query for field '{column}' requested {requested:?} document granularity, \
but the existing FTS index uses a different granularity: {indexed}"
)));
}
None if available.is_empty() => DocumentGranularity::Row,
None if available.len() == 1 => available[0],
None => {
let indexed = indices
.iter()
.map(|(name, document_granularity)| format!("'{name}' ({document_granularity:?})"))
.collect::<Vec<_>>()
.join(", ");
return Err(Error::invalid_input(format!(
"FTS query for field '{column}' is ambiguous because Row and ListElement \
indexes coexist: {indexed}; specify document_granularity"
)));
}
};
resolve_fts_field(dataset.schema(), column, resolved)?;
Ok(resolved)
}
pub(crate) fn fts_document_schema(coordinate_rank: usize) -> Arc<ArrowSchema> {
let mut fields = vec![
ArrowField::new(VALUE_COLUMN_NAME, DataType::Utf8, false),
ArrowField::new(ROW_ID, DataType::UInt64, false),
];
fields.extend(
(0..coordinate_rank)
.map(|rank| ArrowField::new(doc_index_storage_column(rank), DataType::UInt32, false)),
);
Arc::new(ArrowSchema::new(fields))
}
pub(crate) fn transform_fts_document_stream(
input: SendableRecordBatchStream,
resolved: ResolvedFtsField,
) -> Result<SendableRecordBatchStream> {
let input_schema = input.schema();
if (input_schema
.column_with_name(&resolved.root_column)
.is_none()
&& input_schema.column_with_name(VALUE_COLUMN_NAME).is_none())
|| input_schema.column_with_name(ROW_ID).is_none()
{
return Err(Error::internal(
"FTS document input must contain the root source column and _rowid".to_string(),
));
}
let output_schema = fts_document_schema(resolved.coordinate_rank());
let stream_schema = output_schema.clone();
let stream = input.map(move |batch| {
let batch = batch?;
let source = batch
.column_by_name(&resolved.root_column)
.or_else(|| batch.column_by_name(VALUE_COLUMN_NAME))
.expect("FTS document input schema was validated");
let documents = resolved
.documents_from_array(source, batch.num_rows())
.map_err(DataFusionError::from)?;
let input_row_ids = batch[ROW_ID].as_primitive::<arrow_array::types::UInt64Type>();
let texts =
StringArray::from_iter_values(documents.iter().map(|document| document.text.as_str()));
let row_ids = UInt64Array::from_iter_values(
documents
.iter()
.map(|document| input_row_ids.value(document.row_index)),
);
let mut columns = vec![Arc::new(texts) as ArrayRef, Arc::new(row_ids) as ArrayRef];
for rank in 0..resolved.coordinate_rank() {
columns.push(Arc::new(UInt32Array::from_iter_values(
documents.iter().map(|document| document.doc_index[rank]),
)) as ArrayRef);
}
RecordBatch::try_new(stream_schema.clone(), columns).map_err(DataFusionError::from)
});
Ok(Box::pin(RecordBatchStreamAdapter::new(
output_schema,
stream,
)))
}
pub(crate) fn normalize_inverted_details(
index: &IndexMetadata,
mut details: InvertedIndexDetails,
) -> Result<InvertedIndexDetails> {
if !matches!(index.index_version, 0..=3) {
return Err(Error::invalid_input(format!(
"FTS index '{}' has unsupported index_version {}; expected 0, 1, 2, or 3",
index.name, index.index_version
)));
}
if details.posting_format_version.is_none() {
let posting_format_version = match index.index_version {
0 | 1 => 1,
2 => 2,
3 => 3,
_ => unreachable!("index version was validated above"),
};
details.posting_format_version = Some(posting_format_version);
}
Ok(details)
}
fn empty_inverted_update_stream(
dataset: &Dataset,
resolved: &ResolvedFtsField,
) -> Result<SendableRecordBatchStream> {
let field = dataset
.schema()
.field_by_id(resolved.final_field_id)
.ok_or_else(|| {
Error::invalid_input(format!(
"merge_existing_index_segments: field id {} does not exist",
resolved.final_field_id
))
})?;
let schema = if resolved.has_lists() {
fts_document_schema(resolved.coordinate_rank())
} else {
Arc::new(ArrowSchema::new(vec![
ArrowField::new(VALUE_COLUMN_NAME, field.data_type(), true),
ArrowField::new(ROW_ID, arrow_schema::DataType::UInt64, false),
]))
};
Ok(Box::pin(RecordBatchStreamAdapter::new(
schema,
futures::stream::empty(),
)))
}
pub(crate) async fn finalize_segment_files_if_needed(
dataset: &Dataset,
segment: &IndexMetadata,
) -> Result<()> {
let index_dir = dataset.indices_dir().join(segment.uuid.to_string());
let metadata_path = index_dir
.clone()
.join(lance_index::scalar::inverted::METADATA_FILE);
if dataset.object_store.as_ref().exists(&metadata_path).await? {
return Ok(());
}
let store = Arc::new(LanceIndexStore::from_dataset_for_new(
dataset,
&segment.uuid,
)?);
lance_index::scalar::inverted::builder::merge_index_files(
dataset.object_store.as_ref(),
&index_dir,
store,
lance_index::progress::noop_progress(),
)
.await
}
pub(crate) async fn merge_segments(
dataset: &Dataset,
segments: Vec<IndexMetadata>,
) -> Result<IndexMetadata> {
if segments.is_empty() {
return Err(Error::index("No segment metadata was provided".to_string()));
}
let field_id = *segments[0].fields.first().ok_or_else(|| {
Error::invalid_input(format!(
"CreateIndex: segment {} is missing field ids",
segments[0].uuid
))
})?;
let details = match segments[0].index_details.as_ref() {
Some(details_any) => {
let details =
InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| {
Error::io(format!(
"failed to decode InvertedIndexDetails payload: {error}"
))
})?;
normalize_inverted_details(&segments[0], details)?
}
None => normalize_inverted_details(&segments[0], InvertedIndexDetails::default())?,
};
let document_granularity = DocumentGranularity::try_from(details.document_granularity)?;
let resolved = resolve_fts_field_by_id(dataset.schema(), field_id, document_granularity)?;
load_segment_details(dataset, &resolved.canonical_path, &segments).await?;
let mut source_indices = Vec::with_capacity(segments.len());
let mut fragment_bitmap = RoaringBitmap::new();
for segment in &segments {
finalize_segment_files_if_needed(dataset, segment).await?;
fragment_bitmap |= segment.fragment_bitmap.as_ref().cloned().ok_or_else(|| {
Error::invalid_input(format!(
"CreateIndex: segment {} is missing fragment coverage",
segment.uuid
))
})?;
if segment.fields != segments[0].fields {
return Err(Error::invalid_input(format!(
"FTS index {} has inconsistent fields across segments",
segments[0].name
)));
}
let scalar_index = super::open_scalar_index(
dataset,
&resolved.canonical_path,
segment,
&NoOpMetricsCollector,
)
.await?;
let inverted_index = scalar_index
.as_any()
.downcast_ref::<InvertedIndex>()
.ok_or_else(|| {
Error::index(format!(
"merge_existing_index_segments: expected inverted segment {}, got {:?}",
segment.uuid,
scalar_index.index_type()
))
})?;
source_indices.push(Arc::new(inverted_index.clone()));
}
let new_uuid = Uuid::new_v4();
let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?;
let created_index = InvertedIndex::merge_segments(
&source_indices,
empty_inverted_update_stream(dataset, &resolved)?,
&new_store,
None,
lance_index::progress::noop_progress(),
)
.await?;
Ok(IndexMetadata {
uuid: new_uuid,
dataset_version: dataset.manifest.version,
fragment_bitmap: Some(fragment_bitmap),
index_details: Some(Arc::new(created_index.index_details)),
index_version: created_index.index_version as i32,
created_at: Some(chrono::Utc::now()),
base_id: None,
files: Some(index_files_to_table(created_index.files)),
..segments[0].clone()
})
}
pub async fn load_segments(
dataset: &Dataset,
column: &str,
document_granularity: DocumentGranularity,
) -> Result<Option<Vec<IndexMetadata>>> {
let Some(index_meta) = dataset
.load_scalar_index(
lance_index::IndexCriteria::default()
.for_column(column)
.supports_fts()
.with_fts_document_granularity(document_granularity),
)
.await?
else {
return Ok(None);
};
let indices = dataset.load_indices_by_name(&index_meta.name).await?;
if indices.is_empty() {
return Ok(None);
}
let expected_fields = indices[0].fields.clone();
for meta in &indices {
if meta.fields != expected_fields {
return Err(Error::invalid_input(format!(
"FTS index {} has inconsistent fields across segments",
index_meta.name
)));
}
}
Ok(Some(indices))
}
pub(crate) async fn fts_index_fragment_bitmap(
dataset: &Dataset,
column: &str,
document_granularity: DocumentGranularity,
) -> Result<Option<RoaringBitmap>> {
let Some(segments) = load_segments(dataset, column, document_granularity).await? else {
return Ok(None);
};
let fragment_bitmap =
segments
.iter()
.try_fold(RoaringBitmap::new(), |mut coverage, segment| {
coverage |= segment.fragment_bitmap.as_ref()?.clone();
Some(coverage)
});
Ok(fragment_bitmap)
}
pub async fn load_segment_details(
dataset: &Dataset,
column: &str,
segments: &[IndexMetadata],
) -> Result<InvertedIndexDetails> {
let mut expected_details: Option<InvertedIndexDetails> = None;
for meta in segments {
let details_any = fetch_index_details(dataset, column, meta).await?;
let details =
InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|err| {
Error::io(format!(
"failed to decode InvertedIndexDetails payload: {err}"
))
})?;
let details = canonicalize_inverted_index_details(details)?;
match &expected_details {
Some(expected) if !inverted_index_details_semantically_equal(expected, &details) => {
return Err(Error::invalid_input(format!(
"FTS index {} has inconsistent inverted index details across segments",
meta.name
)));
}
Some(_) => {}
None => expected_details = Some(details),
}
}
expected_details.ok_or_else(|| {
Error::invalid_input(format!(
"FTS index for column {} requires at least one segment",
column
))
})
}
fn canonicalize_inverted_index_details(
details: InvertedIndexDetails,
) -> Result<InvertedIndexDetails> {
let params = InvertedIndexParams::try_from(&details)?;
InvertedIndexDetails::try_from(¶ms)
}
fn inverted_index_details_semantically_equal(
left: &InvertedIndexDetails,
right: &InvertedIndexDetails,
) -> bool {
let mut left = left.clone();
let mut right = right.clone();
left.posting_format_version = None;
right.posting_format_version = None;
left == right
}
pub async fn load_segment_params(
dataset: &Dataset,
segment: &IndexMetadata,
) -> Result<InvertedIndexParams> {
let store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?;
InvertedIndex::load_params(&store).await
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_resolve_query_document_granularity_recognizes_covered_index() {
use arrow_array::builder::{ListBuilder, StringBuilder};
use arrow_array::{Int32Array, RecordBatchIterator};
use lance_core::utils::tempfile::TempStrDir;
use lance_index::IndexType;
let mut tags_builder = ListBuilder::new(StringBuilder::new());
for tag in ["alpha", "beta", "gamma", "delta"] {
tags_builder.values().append_value(tag);
tags_builder.append(true);
}
let tags: ArrayRef = Arc::new(tags_builder.finish());
let ids: ArrayRef = Arc::new(Int32Array::from_iter_values(0..4));
let batch = RecordBatch::try_from_iter(vec![("id", ids), ("tags", tags)]).unwrap();
let schema = batch.schema();
let test_dir = TempStrDir::default();
let mut dataset = Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema),
test_dir.as_str(),
None,
)
.await
.unwrap();
let tags_field_id = dataset.schema().field("tags").unwrap().id;
let id_field_id = dataset.schema().field("id").unwrap().id;
let params =
InvertedIndexParams::default().document_granularity(DocumentGranularity::ListElement);
let segment = dataset
.create_index_builder(&["tags"], IndexType::Inverted, ¶ms)
.name("tags_idx".to_string())
.execute_uncommitted()
.await
.unwrap();
let covered = IndexMetadata {
fields: vec![tags_field_id, id_field_id],
covering_fields: vec![id_field_id],
..segment
};
dataset
.commit_existing_index_segments("tags_idx", "tags", vec![covered])
.await
.unwrap();
let resolved = resolve_query_document_granularity(&dataset, "tags", None)
.await
.unwrap();
assert_eq!(
resolved,
DocumentGranularity::ListElement,
"a covered ListElement FTS index must still be recognized, not silently \
skipped in favor of the Row default"
);
}
fn fts_test_schema() -> Schema {
let schema = ArrowSchema::new(vec![
ArrowField::new("text", DataType::Utf8, true),
ArrowField::new(
"tags",
DataType::List(Arc::new(ArrowField::new("item", DataType::Utf8, true))),
true,
),
]);
Schema::try_from(&schema).unwrap()
}
#[test]
fn decode_legacy_inverted_details_type_url() {
let mut details_any = prost_types::Any::from_msg(&InvertedIndexDetails::default()).unwrap();
details_any.type_url = "/lance.index.pb.InvertedIndexDetails".to_string();
let decoded = InvertedIndexDetails::decode(details_any.value.as_slice()).unwrap();
assert_eq!(decoded, InvertedIndexDetails::default());
}
#[test]
fn resolve_element_document_field() {
let schema = fts_test_schema();
let tags = schema.field("tags").unwrap();
let resolved =
resolve_fts_field(&schema, "tags", DocumentGranularity::ListElement).unwrap();
assert_eq!(resolved.final_field_id, tags.id);
assert_eq!(resolved.root_column, "tags");
assert_eq!(resolved.canonical_path, "tags");
assert_eq!(resolved.coordinate_rank(), 1);
let err = resolve_fts_field(&schema, "text", DocumentGranularity::ListElement).unwrap_err();
assert!(err.to_string().contains("has no List layer"), "{err}");
let err =
resolve_fts_field(&schema, "tags[*]", DocumentGranularity::ListElement).unwrap_err();
assert!(err.to_string().contains("does not exist"), "{err}");
}
#[test]
fn normalize_legacy_metadata_as_row_document() {
let metadata = IndexMetadata {
uuid: Uuid::new_v4(),
fields: vec![1],
covering_fields: vec![],
name: "tags_idx".to_string(),
dataset_version: 1,
fragment_bitmap: None,
index_details: None,
index_version: 3,
created_at: None,
base_id: None,
files: None,
};
let details =
normalize_inverted_details(&metadata, InvertedIndexDetails::default()).unwrap();
assert_eq!(
DocumentGranularity::try_from(details.document_granularity).unwrap(),
DocumentGranularity::Row
);
assert_eq!(details.posting_format_version, Some(3));
}
#[test]
fn normalize_rejects_unreleased_v4_metadata() {
let metadata = IndexMetadata {
uuid: Uuid::new_v4(),
fields: vec![1],
covering_fields: vec![],
name: "tags_idx".to_string(),
dataset_version: 1,
fragment_bitmap: None,
index_details: None,
index_version: 4,
created_at: None,
base_id: None,
files: None,
};
let details = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap();
let err = normalize_inverted_details(&metadata, details).unwrap_err();
assert!(err.to_string().contains("unsupported index_version 4"));
}
#[test]
fn canonicalize_inverted_details_accepts_legacy_empty_details() {
let legacy = InvertedIndexDetails::default();
let current = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap();
assert_ne!(legacy, current);
assert_eq!(
canonicalize_inverted_index_details(legacy).unwrap(),
canonicalize_inverted_index_details(current).unwrap()
);
}
#[test]
fn inverted_details_equal_when_only_posting_format_version_differs() {
let left = canonicalize_inverted_index_details(
InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(),
)
.unwrap();
let mut right = left.clone();
right.posting_format_version = Some(1);
assert_ne!(left, right);
assert!(inverted_index_details_semantically_equal(&left, &right));
}
#[test]
fn inverted_details_reject_with_position_mismatch() {
let left = canonicalize_inverted_index_details(
InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(),
)
.unwrap();
let mut right = left.clone();
right.with_position = !left.with_position;
assert!(!inverted_index_details_semantically_equal(&left, &right));
}
}