use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions};
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::union::UnionExec;
use datafusion::prelude::Expr;
use lance_core::{Error, Result, is_system_column};
use lance_index::scalar::FullTextSearchQuery;
use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, Operator};
use tracing::instrument;
use super::block_list::compute_source_block_lists;
use super::collector::LsmDataSourceCollector;
use super::data_source::LsmDataSource;
use super::exec::PkBlockFilterExec;
use super::flushed_cache::{DatasetCache, GenerationWarmer, open_flushed_dataset};
use super::projection::{project_to_canonical, validate_projection_names};
use crate::dataset::mem_wal::memtable::scanner::MemTableScanner;
use crate::session::Session;
use lance_io::object_store::ObjectStoreParams;
pub const SCORE_COLUMN: &str = "_score";
const DEFAULT_OVERFETCH_FACTOR: f64 = 1.0;
fn validate_lsm_fts_query(query: &FullTextSearchQuery) -> Result<()> {
match &query.query {
IndexFtsQuery::Match(m) => {
if m.fuzziness != Some(0) && m.operator != Operator::Or {
return Err(Error::not_supported(
"LSM fuzzy full-text search only supports OR match operators".to_string(),
));
}
Ok(())
}
IndexFtsQuery::Phrase(_) => Ok(()),
_ => Err(Error::not_supported(
"LSM full-text search only supports match and phrase leaf queries".to_string(),
)),
}
}
fn active_source_can_execute_fts(source: &LsmDataSource, column: &str) -> bool {
match source {
LsmDataSource::ActiveMemTable {
batch_store,
index_store,
..
} => {
index_store
.get_fts_by_column(column)
.is_some_and(|index| !index.is_empty())
&& batch_store
.max_visible_row(index_store.max_visible_batch_position())
.is_some()
}
_ => false,
}
}
pub struct LsmFtsSearchPlanner {
collector: LsmDataSourceCollector,
pk_columns: Vec<String>,
base_schema: SchemaRef,
session: Option<Arc<Session>>,
store_params: Option<ObjectStoreParams>,
flushed_cache: Option<Arc<dyn DatasetCache>>,
warmer: Option<Arc<dyn GenerationWarmer>>,
overfetch_factor: f64,
filter: Option<Expr>,
}
impl LsmFtsSearchPlanner {
pub fn new(
collector: LsmDataSourceCollector,
pk_columns: Vec<String>,
base_schema: SchemaRef,
) -> Self {
Self {
collector,
pk_columns,
base_schema,
session: None,
store_params: None,
flushed_cache: None,
warmer: None,
overfetch_factor: DEFAULT_OVERFETCH_FACTOR,
filter: None,
}
}
pub fn with_filter(mut self, filter: Option<Expr>) -> Self {
self.filter = filter;
self
}
pub fn with_overfetch_factor(mut self, factor: f64) -> Self {
self.overfetch_factor = factor;
self
}
pub fn with_session(mut self, session: Arc<Session>) -> Self {
self.session = Some(session);
self
}
pub fn with_store_params(mut self, store_params: ObjectStoreParams) -> Self {
self.store_params = Some(store_params);
self
}
pub fn with_flushed_cache(mut self, cache: Arc<dyn DatasetCache>) -> Self {
self.flushed_cache = Some(cache);
self
}
pub fn with_warmer(mut self, warmer: Arc<dyn GenerationWarmer>) -> Self {
self.warmer = Some(warmer);
self
}
#[instrument(
name = "lsm_fts_search",
level = "info",
skip_all,
fields(column = %column, limit)
)]
pub async fn plan_search(
&self,
column: &str,
query: FullTextSearchQuery,
limit: Option<usize>,
projection: Option<&[String]>,
) -> Result<Arc<dyn ExecutionPlan>> {
let sources = self.collector.collect()?;
if sources
.iter()
.any(|source| active_source_can_execute_fts(source, column))
{
validate_lsm_fts_query(&query)?;
}
validate_projection_names(projection, &self.base_schema, &[SCORE_COLUMN])?;
let target_schema = self.canonical_fts_schema(projection);
let overfetch = super::validate_overfetch_factor(self.overfetch_factor)?;
if sources.is_empty() {
return self.empty_plan(&target_schema);
}
let block_lists = Box::pin(compute_source_block_lists(
&sources,
self.session.as_ref(),
self.store_params.as_ref(),
self.flushed_cache.as_ref(),
))
.await?;
let arm_inputs: Vec<_> = sources
.iter()
.map(|source| {
let is_active = matches!(source, LsmDataSource::ActiveMemTable { .. });
let blocked = block_lists.get(&(source.shard_id(), source.generation()));
let active_needs_uncapped_recency = match source {
LsmDataSource::ActiveMemTable { index_store, .. }
if !self.pk_columns.is_empty() =>
{
!index_store.has_pk_index() || index_store.pk_has_overrides()
}
_ => false,
};
let fetch_limit = if active_needs_uncapped_recency {
None
} else if blocked.is_some() && !self.pk_columns.is_empty() {
limit.map(|limit| ((limit as f64) * overfetch).ceil() as usize)
} else {
limit
};
(source, is_active, blocked, fetch_limit)
})
.collect();
let built =
futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, fetch_limit)| {
Box::pin(self.build_source_plan(source, column, &query, *fetch_limit, projection))
}))
.await?;
let mut per_source_plans: Vec<Arc<dyn ExecutionPlan>> = Vec::with_capacity(sources.len());
for ((_, _, blocked, _), plan) in arm_inputs.iter().zip(built) {
let blocked = *blocked;
let deduped = if let Some(set) = blocked
&& !self.pk_columns.is_empty()
{
Arc::new(PkBlockFilterExec::new(
plan,
self.pk_columns.clone(),
set.clone(),
limit.unwrap_or(usize::MAX),
)) as Arc<dyn ExecutionPlan>
} else {
plan
};
let normalized = project_to_canonical(deduped, &target_schema)?;
per_source_plans.push(normalized);
}
let merged: Arc<dyn ExecutionPlan> = if per_source_plans.len() == 1 {
per_source_plans.into_iter().next().unwrap()
} else {
#[allow(deprecated)]
Arc::new(UnionExec::new(per_source_plans))
};
let score_idx = merged.schema().index_of(SCORE_COLUMN).map_err(|_| {
Error::internal(format!(
"{SCORE_COLUMN} missing from canonical FTS schema after merge"
))
})?;
let sort_expr = vec![PhysicalSortExpr {
expr: Arc::new(Column::new(SCORE_COLUMN, score_idx)),
options: SortOptions {
descending: true,
nulls_first: false,
},
}];
let lex_ordering = LexOrdering::new(sort_expr).ok_or_else(|| {
Error::internal("Failed to build LexOrdering for FTS _score sort".to_string())
})?;
let per_partition_sorted: Arc<dyn ExecutionPlan> = Arc::new(
SortExec::new(lex_ordering.clone(), merged)
.with_preserve_partitioning(true)
.with_fetch(limit),
);
let merged_sorted: Arc<dyn ExecutionPlan> = Arc::new(
SortPreservingMergeExec::new(lex_ordering, per_partition_sorted).with_fetch(limit),
);
Ok(merged_sorted)
}
async fn build_source_plan(
&self,
source: &LsmDataSource,
column: &str,
query: &FullTextSearchQuery,
limit: Option<usize>,
projection: Option<&[String]>,
) -> Result<Arc<dyn ExecutionPlan>> {
match source {
LsmDataSource::BaseTable { dataset } => {
let mut scanner = dataset.scan();
let cols = self.fts_scanner_projection(projection);
scanner.project(&cols.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
if let Some(ref filter) = self.filter {
scanner.filter_expr(filter.clone());
scanner.prefilter(true);
}
let mut bound_query = query.clone().with_column(column.to_string())?;
if let Some(limit) = limit {
bound_query = bound_query.limit(Some(limit as i64));
} else {
bound_query = bound_query.limit(None);
}
scanner.full_text_search(bound_query)?;
scanner.create_plan().await
}
LsmDataSource::FlushedMemTable { path, .. } => {
let dataset = open_flushed_dataset(
path,
self.session.as_ref(),
self.store_params.as_ref(),
self.flushed_cache.as_ref(),
self.warmer.as_ref(),
)
.await?;
let mut scanner = dataset.scan();
let cols = self.fts_scanner_projection(projection);
scanner.project(&cols.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
if let Some(ref filter) = self.filter {
scanner.filter_expr(filter.clone());
scanner.prefilter(true);
}
let mut bound_query = query.clone().with_column(column.to_string())?;
if let Some(limit) = limit {
bound_query = bound_query.limit(Some(limit as i64));
} else {
bound_query = bound_query.limit(None);
}
scanner.full_text_search(bound_query)?;
scanner.create_plan().await
}
LsmDataSource::ActiveMemTable {
batch_store,
index_store,
schema,
..
} => {
if !active_source_can_execute_fts(source, column) {
return self.empty_plan(&self.canonical_fts_schema(projection));
}
validate_lsm_fts_query(query)?;
let mut scanner =
MemTableScanner::new(batch_store.clone(), index_store.clone(), schema.clone());
let cols = self.fts_scanner_projection(projection);
scanner.project(&cols.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
if let Some(ref filter) = self.filter {
scanner.filter_expr(filter.clone());
}
if !self.pk_columns.is_empty() {
scanner.with_pk_columns(self.pk_columns.clone());
}
let mut bound_query = query.clone().with_column(column.to_string())?;
if let Some(limit) = limit {
bound_query = bound_query.limit(Some(limit as i64));
} else {
bound_query = bound_query.limit(None);
}
scanner.full_text_search(bound_query)?;
scanner.create_plan().await
}
}
}
fn fts_scanner_projection(&self, user_projection: Option<&[String]>) -> Vec<String> {
let mut cols: Vec<String> = if let Some(p) = user_projection {
p.iter()
.filter(|c| !is_system_column(c) && c.as_str() != SCORE_COLUMN)
.cloned()
.collect()
} else {
self.base_schema
.fields()
.iter()
.map(|f| f.name().clone())
.collect()
};
for pk in &self.pk_columns {
if !cols.contains(pk) {
cols.push(pk.clone());
}
}
cols
}
fn canonical_fts_schema(&self, user_projection: Option<&[String]>) -> SchemaRef {
let mut ordered: Vec<String> = if let Some(p) = user_projection {
p.to_vec()
} else {
self.base_schema
.fields()
.iter()
.map(|f| f.name().clone())
.collect()
};
for pk in &self.pk_columns {
if !ordered.contains(pk) {
ordered.push(pk.clone());
}
}
if !ordered.iter().any(|c| c == SCORE_COLUMN) {
ordered.push(SCORE_COLUMN.to_string());
}
let fields: Vec<Arc<Field>> = ordered
.iter()
.filter_map(|name| {
if name == SCORE_COLUMN {
Some(Arc::new(Field::new(SCORE_COLUMN, DataType::Float32, true)))
} else if is_system_column(name) {
Some(Arc::new(Field::new(name.clone(), DataType::UInt64, true)))
} else {
self.base_schema
.field_with_name(name)
.ok()
.map(|f| Arc::new(f.clone()))
}
})
.collect();
Arc::new(Schema::new(fields))
}
fn empty_plan(&self, schema: &SchemaRef) -> Result<Arc<dyn ExecutionPlan>> {
use datafusion::physical_plan::empty::EmptyExec;
Ok(Arc::new(EmptyExec::new(schema.clone())))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables};
use crate::dataset::mem_wal::write::{BatchStore, IndexStore};
use crate::dataset::{Dataset, WriteParams};
use arrow_array::{BooleanArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray};
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
use futures::TryStreamExt;
use std::collections::HashMap;
fn fts_schema() -> Arc<ArrowSchema> {
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let id_field = Field::new("id", DataType::Int32, false).with_metadata(id_meta);
Arc::new(ArrowSchema::new(vec![
id_field,
Field::new("text", DataType::Utf8, true),
]))
}
fn fts_tombstone_schema() -> Arc<ArrowSchema> {
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let id_field = Field::new("id", DataType::Int32, false).with_metadata(id_meta);
Arc::new(ArrowSchema::new(vec![
id_field,
Field::new("text", DataType::Utf8, true),
Field::new(crate::dataset::mem_wal::TOMBSTONE, DataType::Boolean, false),
]))
}
fn make_batch(schema: &ArrowSchema, ids: &[i32], texts: &[&str]) -> RecordBatch {
RecordBatch::try_new(
Arc::new(schema.clone()),
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(StringArray::from(texts.to_vec())),
],
)
.unwrap()
}
fn make_tombstone_batch(
schema: &ArrowSchema,
rows: &[(i32, Option<&str>, bool)],
) -> RecordBatch {
let ids: Vec<i32> = rows.iter().map(|(id, _, _)| *id).collect();
let texts: Vec<Option<&str>> = rows.iter().map(|(_, text, _)| *text).collect();
let tombstones: Vec<bool> = rows.iter().map(|(_, _, tombstone)| *tombstone).collect();
RecordBatch::try_new(
Arc::new(schema.clone()),
vec![
Arc::new(Int32Array::from(ids)),
Arc::new(StringArray::from(texts)),
Arc::new(BooleanArray::from(tombstones)),
],
)
.unwrap()
}
async fn write_dataset(uri: &str, batches: Vec<RecordBatch>) -> Dataset {
let schema = batches[0].schema();
let has_id = schema.column_with_name("id").is_some();
let reader = RecordBatchIterator::new(batches.clone().into_iter().map(Ok), schema);
let dataset = Dataset::write(reader, uri, Some(WriteParams::default()))
.await
.unwrap();
if has_id {
crate::dataset::mem_wal::scanner::block_list::write_pk_sidecar(uri, &batches, &["id"])
.await
.unwrap();
}
dataset
}
#[tokio::test]
async fn rejects_missing_projection_column() {
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let projection = vec!["missing".to_string()];
let err = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
Some(&projection),
)
.await
.unwrap_err();
assert!(
err.to_string().contains("missing"),
"unexpected missing-column projection error: {err}"
);
}
#[tokio::test]
async fn local_mode_unions_base_and_active_with_consistent_score_schema() {
use crate::index::DatasetIndexExt;
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_batch(
&schema,
&[1, 2],
&["lance rocks", "unrelated text"],
)],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_batch(
&schema,
&[3, 4],
&["lance memwal goes fast", "completely unrelated"],
);
batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a base+active union plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total: usize = batches.iter().map(|b| b.num_rows()).sum();
assert!(
total >= 2,
"expected at least the 2 'lance' rows from base+active, got {total}"
);
let out = batches[0].schema();
let score_field = out
.field_with_name(SCORE_COLUMN)
.expect("_score column missing from output");
assert!(
score_field.is_nullable(),
"_score must be nullable to stay union-compatible across base+active"
);
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert!(ids.contains(&1), "missing base hit id=1; got ids={ids:?}");
assert!(ids.contains(&3), "missing active hit id=3; got ids={ids:?}");
}
#[tokio::test]
async fn prefilter_drops_nonmatching_hits_across_base_and_active() {
use crate::index::DatasetIndexExt;
use datafusion::prelude::{col, lit};
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_batch(&schema, &[1, 2], &["lance one", "lance two"])],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_batch(&schema, &[0, 3], &["lance zero", "lance three"]);
batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("id").gt_eq(lit(2i32))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a filtered base+active plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
ids.sort();
assert_eq!(
ids,
vec![2, 3],
"prefilter must keep only id>=2 across base (id=2) and active (id=3), \
dropping base id=1 and active id=0; got {ids:?}"
);
}
#[tokio::test]
async fn prefilter_on_flushed_composes_with_block_list() {
use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot;
use crate::index::DatasetIndexExt;
use datafusion::prelude::{col, lit};
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let shard_id = uuid::Uuid::new_v4();
let gen1_uri = format!("{}/_mem_wal/{}/gen_1", base_uri, shard_id);
let mut gen1 = write_dataset(
&gen1_uri,
vec![make_batch(
&schema,
&[1, 3, 4],
&["lance lance lance lance", "lance lance lance", "lance"],
)],
)
.await;
gen1.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let gen2_uri = format!("{}/_mem_wal/{}/gen_2", base_uri, shard_id);
let mut gen2 =
write_dataset(&gen2_uri, vec![make_batch(&schema, &[3], &["other text"])]).await;
gen2.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let snapshot = ShardSnapshot::new(shard_id)
.with_current_generation(3)
.with_flushed_generation(1, "gen_1".to_string())
.with_flushed_generation(2, "gen_2".to_string());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![snapshot]);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("id").gt_eq(lit(3i32))))
.with_overfetch_factor(2.0);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
None,
)
.await
.expect("planner should produce a filtered flushed plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert_eq!(
ids,
vec![4],
"flushed FTS prefilter should return live id=4 after stale id=3 is blocked; got {ids:?}"
);
}
#[tokio::test]
async fn active_tombstone_masks_base_fts_hit() {
use crate::index::DatasetIndexExt;
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let base_schema = fts_schema();
let mem_schema = fts_tombstone_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base = write_dataset(
&base_uri,
vec![make_batch(&base_schema, &[1, 2], &["lance lance", "lance"])],
)
.await;
base.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let active_tombstone = make_tombstone_batch(&mem_schema, &[(1, None, true)]);
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut index_store = IndexStore::new();
index_store.enable_pk_index(&[("id".to_string(), 0)]);
index_store.add_fts("text_fts".to_string(), 1, "text".to_string());
let (_, row_offset, batch_position) = batch_store.append(active_tombstone.clone()).unwrap();
index_store
.insert_with_batch_position(&active_tombstone, row_offset, Some(batch_position))
.unwrap();
let index_store = Arc::new(index_store);
let collector = LsmDataSourceCollector::new(Arc::new(base), vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store,
schema: mem_schema,
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], base_schema)
.with_overfetch_factor(2.0);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
None,
)
.await
.unwrap();
let stream = plan
.execute(0, datafusion::prelude::SessionContext::new().task_ctx())
.unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert_eq!(
ids,
vec![2],
"active tombstone for id=1 must block the older base FTS hit; got {ids:?}"
);
}
#[tokio::test]
async fn active_update_and_tombstone_mask_frozen_fts_hits() {
let base_schema = fts_schema();
let mem_schema = fts_tombstone_schema();
let frozen_batch = make_tombstone_batch(
&mem_schema,
&[
(1, Some("lance stale update"), false),
(2, Some("lance live"), false),
(3, Some("lance deleted"), false),
],
);
let frozen_batch_store = Arc::new(BatchStore::with_capacity(16));
let mut frozen_index_store = IndexStore::new();
frozen_index_store.enable_pk_index(&[("id".to_string(), 0)]);
frozen_index_store.add_fts("text_fts".to_string(), 1, "text".to_string());
let (_, frozen_row_offset, frozen_batch_position) =
frozen_batch_store.append(frozen_batch.clone()).unwrap();
frozen_index_store
.insert_with_batch_position(
&frozen_batch,
frozen_row_offset,
Some(frozen_batch_position),
)
.unwrap();
let active_batch = make_tombstone_batch(
&mem_schema,
&[(1, Some("fresh other text"), false), (3, None, true)],
);
let active_batch_store = Arc::new(BatchStore::with_capacity(16));
let mut active_index_store = IndexStore::new();
active_index_store.enable_pk_index(&[("id".to_string(), 0)]);
active_index_store.add_fts("text_fts".to_string(), 1, "text".to_string());
let (_, active_row_offset, active_batch_position) =
active_batch_store.append(active_batch.clone()).unwrap();
active_index_store
.insert_with_batch_position(
&active_batch,
active_row_offset,
Some(active_batch_position),
)
.unwrap();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let shard_id = uuid::Uuid::new_v4();
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
shard_id,
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store: active_batch_store,
index_store: Arc::new(active_index_store),
schema: mem_schema.clone(),
generation: 2,
},
frozen: vec![InMemoryMemTableRef {
batch_store: frozen_batch_store,
index_store: Arc::new(frozen_index_store),
schema: mem_schema,
generation: 1,
}],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], base_schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.unwrap();
let stream = plan
.execute(0, datafusion::prelude::SessionContext::new().task_ctx())
.unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
ids.sort_unstable();
assert_eq!(
ids,
vec![2],
"active update id=1 and tombstone id=3 must block stale frozen FTS hits; got {ids:?}"
);
}
#[tokio::test]
async fn active_filtered_search_without_pk_applies_small_limit_after_filter() {
use datafusion::prelude::{col, lit};
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_batch(
&schema,
&[1, 2, 3],
&["lance", "lance filler", "lance filler filler"],
);
batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec![], schema)
.with_filter(Some(col("id").gt_eq(lit(1i32))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(2),
None,
)
.await
.expect("planner should produce an active-only filtered plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
ids.sort_unstable();
assert_eq!(
ids,
vec![1, 2],
"no-PK filtered active search must apply the limit after filtering \
and keep the top-scoring matching hits; got ids={ids:?}"
);
}
#[tokio::test]
async fn fuzzy_and_query_is_rejected_when_active_memtable_is_present() {
use lance_index::scalar::inverted::query::{
FtsQuery as IndexFtsQuery, MatchQuery, Operator,
};
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_batch(&schema, &[1], &["lance memwal"]);
let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, row_offset, Some(batch_position))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec![], schema);
let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match(
MatchQuery::new("lance memwal".to_string())
.with_operator(Operator::And)
.with_fuzziness(Some(1)),
));
let err = planner
.plan_search("text", query, Some(10), None)
.await
.expect_err("fuzzy AND should be rejected consistently");
assert!(
err.to_string().contains("fuzzy full-text search"),
"unexpected error for fuzzy AND query: {err}"
);
}
#[tokio::test]
async fn base_only_boolean_query_uses_dataset_scanner_support() {
use crate::index::DatasetIndexExt;
use lance_index::IndexType;
use lance_index::scalar::inverted::query::{
BooleanQuery, FtsQuery as IndexFtsQuery, MatchQuery, Occur,
};
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_batch(
&schema,
&[1, 2],
&["lance rocks", "unrelated text"],
)],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let collector = LsmDataSourceCollector::new(base_ds, vec![]);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let query = FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new(
vec![(Occur::Must, MatchQuery::new("lance".to_string()).into())],
)));
let plan = planner
.plan_search("text", query, Some(10), Some(&["id".to_string()]))
.await
.expect("base-only boolean query should be delegated to dataset scanner");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let ids: Vec<i32> = batches
.iter()
.flat_map(|batch| {
batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.values()
.to_vec()
})
.collect();
assert_eq!(ids, vec![1]);
}
#[tokio::test]
async fn boolean_query_ignores_active_memtable_without_relevant_fts_index() {
use crate::index::DatasetIndexExt;
use lance_index::IndexType;
use lance_index::scalar::inverted::query::{
BooleanQuery, FtsQuery as IndexFtsQuery, MatchQuery, Occur,
};
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_batch(
&schema,
&[1, 2],
&["lance rocks", "unrelated text"],
)],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
let active_batch = make_batch(&schema, &[99], &["active text has no fts index"]);
let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, row_offset, Some(batch_position))
.unwrap();
let indexes = Arc::new(indexes);
let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let query = FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new(
vec![(Occur::Must, MatchQuery::new("lance".to_string()).into())],
)));
let plan = planner
.plan_search("text", query, Some(10), Some(&["id".to_string()]))
.await
.expect("irrelevant active memtable must not reject base-supported boolean query");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let ids: Vec<i32> = batches
.iter()
.flat_map(|batch| {
batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.values()
.to_vec()
})
.collect();
assert_eq!(ids, vec![1]);
}
#[tokio::test]
async fn prefilter_on_base_is_not_a_lossy_postfilter() {
use crate::index::DatasetIndexExt;
use datafusion::prelude::{col, lit};
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let schema = fts_schema();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_batch(
&schema,
&[1, 2],
&["lance", "lance filler filler filler filler filler"],
)],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let collector = LsmDataSourceCollector::new(base_ds, vec![]);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("id").gt_eq(lit(2i32))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
None,
)
.await
.expect("planner should produce a filtered base plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert_eq!(
ids,
vec![2],
"prefilter must return the lower-scoring match id=2, not post-filter \
the top-1 (id=1) down to nothing; got {ids:?}"
);
}
#[tokio::test]
async fn prefilter_on_active_is_not_a_lossy_postfilter() {
use datafusion::prelude::{col, lit};
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let schema = Arc::new(ArrowSchema::new(vec![
Field::new("id", DataType::Int32, false).with_metadata(id_meta),
Field::new("text", DataType::Utf8, true),
Field::new("status", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec![
"lance",
"lance filler filler filler filler filler",
])),
Arc::new(StringArray::from(vec!["archived", "active"])),
],
)
.unwrap();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
batch_store.append(batch.clone()).unwrap();
indexes
.insert_with_batch_position(&batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("status").eq(lit("active"))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
None,
)
.await
.expect("planner should produce a filtered active plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert_eq!(
ids,
vec![2],
"active FTS prefilter must return the lower-scoring matching row, \
not post-filter the top-1 down to nothing; got {ids:?}"
);
}
#[tokio::test]
async fn active_limit_applies_after_newest_pk_recency_filter() {
use datafusion::prelude::{col, lit};
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let schema = Arc::new(ArrowSchema::new(vec![
Field::new("id", DataType::Int32, false).with_metadata(id_meta),
Field::new("text", DataType::Utf8, true),
Field::new("status", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 1])),
Arc::new(StringArray::from(vec![
"lance",
"lance filler filler filler filler filler",
"other text",
])),
Arc::new(StringArray::from(vec!["active", "active", "active"])),
],
)
.unwrap();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
batch_store.append(batch.clone()).unwrap();
indexes
.insert_with_batch_position(&batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("status").eq(lit("active"))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(1),
None,
)
.await
.expect("planner should produce a filtered active plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert_eq!(
ids,
vec![2],
"active FTS limit must apply after newest-PK filtering; got {ids:?}"
);
}
#[tokio::test]
async fn prefilter_excludes_pk_whose_newest_version_fails() {
use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables};
use crate::dataset::mem_wal::write::{BatchStore, IndexStore};
use arrow_schema::{DataType, Field};
use datafusion::prelude::{col, lit};
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let schema = Arc::new(ArrowSchema::new(vec![
Field::new("id", DataType::Int32, false).with_metadata(id_meta),
Field::new("text", DataType::Utf8, true),
Field::new("status", DataType::Utf8, false),
]));
let make_row = |statuses: &[&str]| -> RecordBatch {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![5; statuses.len()])),
Arc::new(StringArray::from(vec!["lance text"; statuses.len()])),
Arc::new(StringArray::from(statuses.to_vec())),
],
)
.unwrap()
};
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let base_batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![999])),
Arc::new(StringArray::from(vec!["unrelated"])),
Arc::new(StringArray::from(vec!["active"])),
],
)
.unwrap();
let base_ds = Arc::new(write_dataset(&base_uri, vec![base_batch]).await);
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_row(&["active", "archived"]);
batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("status").eq(lit("active"))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a filtered active plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(
total, 0,
"pk=5's current version is 'archived' and must be excluded; the stale \
'active' older hit must not leak (filter evaluated on newest version)"
);
}
#[tokio::test]
async fn prefilter_blocks_base_hit_when_active_newest_fails() {
use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables};
use crate::dataset::mem_wal::write::{BatchStore, IndexStore};
use crate::index::DatasetIndexExt;
use datafusion::prelude::{col, lit};
use lance_index::IndexType;
use lance_index::scalar::inverted::tokenizer::InvertedIndexParams;
let mut id_meta = HashMap::new();
id_meta.insert(
"lance-schema:unenforced-primary-key".to_string(),
"true".to_string(),
);
let schema = Arc::new(ArrowSchema::new(vec![
Field::new("id", DataType::Int32, false).with_metadata(id_meta),
Field::new("text", DataType::Utf8, true),
Field::new("status", DataType::Utf8, false),
]));
let make_rows = |rows: &[(i32, &str, &str)]| -> RecordBatch {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(
rows.iter().map(|(id, _, _)| *id).collect::<Vec<_>>(),
)),
Arc::new(StringArray::from(
rows.iter().map(|(_, text, _)| *text).collect::<Vec<_>>(),
)),
Arc::new(StringArray::from(
rows.iter()
.map(|(_, _, status)| *status)
.collect::<Vec<_>>(),
)),
],
)
.unwrap()
};
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let mut base_ds = write_dataset(
&base_uri,
vec![make_rows(&[
(5, "lance base stale", "active"),
(6, "lance base live", "active"),
])],
)
.await;
base_ds
.create_index(
&["text"],
IndexType::Inverted,
Some("text_fts".to_string()),
&InvertedIndexParams::default(),
false,
)
.await
.unwrap();
let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap());
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let active_batch = make_rows(&[(5, "lance active newest", "archived")]);
batch_store.append(active_batch.clone()).unwrap();
indexes
.insert_with_batch_position(&active_batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema)
.with_filter(Some(col("status").eq(lit("active"))));
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a filtered base+active plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids = Vec::new();
for batch in &batches {
let id_array = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for row in 0..batch.num_rows() {
ids.push(id_array.value(row));
}
}
ids.sort_unstable();
assert_eq!(
ids,
vec![6],
"base pk=5 passes the filter but is superseded by active archived pk=5; got {ids:?}"
);
}
#[tokio::test]
async fn local_mode_active_memtable_only_returns_score_sorted_hits() {
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let batch = make_batch(
&schema,
&[1, 2, 3, 4],
&[
"lance is a columnar data format",
"memwal handles streaming writes",
"lance memwal lance lance",
"completely unrelated",
],
);
batch_store.append(batch.clone()).unwrap();
indexes
.insert_with_batch_position(&batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("local mode planner should produce a plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total: usize = batches.iter().map(|b| b.num_rows()).sum();
assert!(
total >= 2,
"expected at least the 2 'lance' rows, got {total}"
);
let out = batches[0].schema();
assert!(out.field_with_name(SCORE_COLUMN).is_ok());
assert!(out.field_with_name("id").is_ok());
let mut prev_score: Option<f32> = None;
for batch in &batches {
let score = batch
.column_by_name(SCORE_COLUMN)
.unwrap()
.as_any()
.downcast_ref::<arrow_array::Float32Array>()
.unwrap();
for i in 0..batch.num_rows() {
let s = score.value(i);
if let Some(p) = prev_score {
assert!(p >= s, "scores not sorted DESC: {p} then {s}");
}
prev_score = Some(s);
}
}
}
#[tokio::test]
async fn active_match_query_preserves_and_operator() {
use lance_index::scalar::inverted::query::{
FtsQuery as IndexFtsQuery, MatchQuery, Operator,
};
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let batch = make_batch(
&schema,
&[1, 2, 3],
&["lance only", "memwal only", "lance memwal"],
);
batch_store.append(batch.clone()).unwrap();
indexes
.insert_with_batch_position(&batch, 0, Some(0))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match(
MatchQuery::new("lance memwal".to_string())
.with_operator(Operator::And)
.with_column(Some("text".to_string())),
));
let plan = planner
.plan_search("text", query, Some(10), None)
.await
.expect("planner should produce an active-only plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
ids.sort_unstable();
assert_eq!(
ids,
vec![3],
"AND query must only return rows containing both terms; got ids={ids:?}"
);
}
#[tokio::test]
async fn local_mode_active_dedups_updated_pk_keeping_newest() {
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let batch_old = make_batch(&schema, &[1, 2], &["lance stale version", "other doc"]);
batch_store.append(batch_old.clone()).unwrap();
indexes
.insert_with_batch_position(&batch_old, 0, Some(0))
.unwrap();
let batch_new = make_batch(&schema, &[1], &["lance fresh version"]);
batch_store.append(batch_new.clone()).unwrap();
indexes
.insert_with_batch_position(&batch_new, 2, Some(1))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("lance".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce an active-only plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut rows: Vec<(i32, String)> = Vec::new();
for b in &batches {
let ids = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
let texts = b
.column_by_name("text")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
for i in 0..b.num_rows() {
rows.push((ids.value(i), texts.value(i).to_string()));
}
}
let id1: Vec<&(i32, String)> = rows.iter().filter(|(id, _)| *id == 1).collect();
assert_eq!(
id1.len(),
1,
"updated PK id=1 must be deduped to one row; got {rows:?}"
);
assert_eq!(
id1[0].1, "lance fresh version",
"dedup must keep the newest (max row-position) version"
);
}
#[tokio::test]
async fn active_stale_update_predicate_crossing_leaks() {
let schema = fts_schema();
let batch_store = Arc::new(BatchStore::with_capacity(16));
let mut indexes = IndexStore::new();
indexes.enable_pk_index(&[("id".to_string(), 0)]);
indexes.add_fts("text_fts".to_string(), 1, "text".to_string());
let b1 = make_batch(&schema, &[1, 2], &["alpha lance", "alpha foo"]);
let (bp1, off1, _) = batch_store.append(b1.clone()).unwrap();
indexes
.insert_with_batch_position(&b1, off1, Some(bp1))
.unwrap();
let b2 = make_batch(&schema, &[1], &["beta lance"]);
let (bp2, off2, _) = batch_store.append(b2.clone()).unwrap();
indexes
.insert_with_batch_position(&b2, off2, Some(bp2))
.unwrap();
let indexes = Arc::new(indexes);
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store,
index_store: indexes,
schema: schema.clone(),
generation: 1,
},
frozen: vec![],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("alpha".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert!(
!ids.contains(&1),
"stale pk=1 (now 'beta lance') leaked on an 'alpha' search; got ids={ids:?}"
);
assert!(
ids.contains(&2),
"live pk=2 ('alpha foo') must still match 'alpha'; got ids={ids:?}"
);
}
#[tokio::test]
async fn cross_gen_stale_update_blocked_by_newer_memtable() {
let schema = fts_schema();
let frozen_store = Arc::new(BatchStore::with_capacity(16));
let mut frozen_idx = IndexStore::new();
frozen_idx.enable_pk_index(&[("id".to_string(), 0)]);
frozen_idx.add_fts("text_fts".to_string(), 1, "text".to_string());
let fb = make_batch(&schema, &[1, 2], &["alpha lance", "alpha foo"]);
let (bp, off, _) = frozen_store.append(fb.clone()).unwrap();
frozen_idx
.insert_with_batch_position(&fb, off, Some(bp))
.unwrap();
let active_store = Arc::new(BatchStore::with_capacity(16));
let mut active_idx = IndexStore::new();
active_idx.enable_pk_index(&[("id".to_string(), 0)]);
active_idx.add_fts("text_fts".to_string(), 1, "text".to_string());
let ab = make_batch(&schema, &[1], &["beta lance"]);
let (bp, off, _) = active_store.append(ab.clone()).unwrap();
active_idx
.insert_with_batch_position(&ab, off, Some(bp))
.unwrap();
let tmp = tempfile::tempdir().unwrap();
let base_uri = format!("{}/base", tmp.path().to_str().unwrap());
let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![])
.with_in_memory_memtables(
uuid::Uuid::new_v4(),
InMemoryMemTables {
active: InMemoryMemTableRef {
batch_store: active_store,
index_store: Arc::new(active_idx),
schema: schema.clone(),
generation: 2,
},
frozen: vec![InMemoryMemTableRef {
batch_store: frozen_store,
index_store: Arc::new(frozen_idx),
schema: schema.clone(),
generation: 1,
}],
},
);
let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema);
let plan = planner
.plan_search(
"text",
FullTextSearchQuery::new("alpha".to_string()),
Some(10),
None,
)
.await
.expect("planner should produce a plan");
let ctx = datafusion::prelude::SessionContext::new();
let stream = plan.execute(0, ctx.task_ctx()).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut ids: Vec<i32> = Vec::new();
for b in &batches {
let col = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..b.num_rows() {
ids.push(col.value(i));
}
}
assert!(
!ids.contains(&1),
"stale frozen pk=1 ('alpha lance', now 'beta lance' in the active gen) \
leaked on an 'alpha' search; got ids={ids:?}"
);
assert!(
ids.contains(&2),
"live pk=2 ('alpha foo', only in the frozen gen) must still match; got ids={ids:?}"
);
}
}