diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 38b60da..8f6d8f4 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -15,7 +15,9 @@ use arrow_array::{ use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, TimeUnit}; use chrono::{DateTime, Timelike, Utc}; use futures::TryStreamExt; -use lance::dataset::mem_wal::{DatasetMemWalExt, ShardWriterConfig}; +use lance::dataset::mem_wal::{ + DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, +}; use lance::dataset::optimize::{compact_files, CompactionMetrics, CompactionOptions}; use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; use lance::index::DatasetIndexExt; @@ -34,6 +36,7 @@ use crate::record::{ContextRecord, SearchResult, StateMetadata}; /// Embedding length used for the semantic index column. const DEFAULT_EMBEDDING_DIM: i32 = 1536; const DEFAULT_SEARCH_LIMIT: usize = 10; +const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; const ID_INDEX_NAME: &str = "id_idx"; /// Configuration for background compaction. @@ -280,18 +283,19 @@ impl ContextStore { limit: Option, offset: Option, ) -> LanceResult> { - let mut scanner = self.dataset.scan(); - if let Some(limit) = limit { - scanner.limit(Some(limit as i64), offset.map(|o| o as i64))?; - } else if let Some(offset) = offset { - scanner.limit(None, Some(offset as i64))?; - } - + let scanner = self.lsm_scanner().await?; let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { results.extend(batch_to_records(&batch)?); } + + if let Some(offset) = offset { + results = results.into_iter().skip(offset).collect(); + } + if let Some(limit) = limit { + results.truncate(limit); + } Ok(results) } @@ -315,18 +319,51 @@ impl ContextStore { return Ok(Vec::new()); } - let query_array = Float32Array::from(query.to_vec()); + let mut results: Vec = self + .list(None, None) + .await? + .into_iter() + .filter_map(|record| { + let distance = l2_distance(query, record.embedding.as_ref()?); + Some(SearchResult { record, distance }) + }) + .collect(); + results.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + results.truncate(top_k); + Ok(results) + } - let mut scanner = self.dataset.scan(); - scanner.nearest("embedding", &query_array, top_k)?; - scanner.limit(Some(top_k as i64), Some(0))?; + async fn lsm_scanner(&self) -> LanceResult { + let object_store = self.dataset.object_store(None).await?; + let branch_location = self.dataset.branch_location(); + let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; + + let mut shard_snapshots = Vec::with_capacity(shard_ids.len()); + for shard_id in shard_ids { + let manifest_store = ShardManifestStore::new( + object_store.clone(), + &branch_location.path, + shard_id, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + let Some(manifest) = manifest_store.read_latest().await? else { + continue; + }; - let mut stream = scanner.try_into_stream().await?; - let mut results = Vec::new(); - while let Some(batch) = stream.try_next().await? { - results.extend(batch_to_search_results(&batch)?); + let mut snapshot = ShardSnapshot::new(shard_id) + .with_spec_id(manifest.shard_spec_id) + .with_current_generation(manifest.current_generation); + for flushed in manifest.flushed_generations { + snapshot = snapshot.with_flushed_generation(flushed.generation, flushed.path); + } + shard_snapshots.push(snapshot); } - Ok(results) + + Ok(LsmScanner::new( + Arc::new(self.dataset.clone()), + shard_snapshots, + vec!["id".to_string()], + )) } /// Manually trigger compaction to merge small fragments. @@ -837,34 +874,6 @@ impl Drop for ContextStore { } } -fn batch_to_search_results(batch: &RecordBatch) -> LanceResult> { - let records = batch_to_records(batch)?; - - let distance_column = batch.column_by_name("_distance").ok_or_else(|| { - LanceError::from(ArrowError::InvalidArgumentError( - "search results missing _distance column".to_string(), - )) - })?; - let distance_array = distance_column - .as_ref() - .as_any() - .downcast_ref::() - .ok_or_else(|| { - LanceError::from(ArrowError::InvalidArgumentError( - "_distance column has unexpected data type".to_string(), - )) - })?; - - Ok(records - .into_iter() - .enumerate() - .map(|(i, record)| SearchResult { - record, - distance: distance_array.value(i), - }) - .collect()) -} - /// Convert a record batch to context records. fn batch_to_records(batch: &RecordBatch) -> LanceResult> { let id_array = column_as::(batch, "id")?; @@ -1071,6 +1080,17 @@ fn embedding_from_list(list: &FixedSizeListArray, row: usize) -> LanceResult f32 { + left.iter() + .zip(right) + .map(|(left, right)| { + let delta = left - right; + delta * delta + }) + .sum::() + .sqrt() +} + fn column_as<'a, A>(batch: &'a RecordBatch, name: &str) -> LanceResult<&'a A> where A: Array + 'static, @@ -1144,16 +1164,15 @@ mod tests { store.add(&[first.clone(), second.clone()]).await.unwrap(); let query = make_embedding(1.0); - let _results = store.search(&query, Some(2)).await.unwrap(); - - // TODO: MemWAL reads are not yet visible via standard scan. - // assert_eq!(results.len(), 2); - // assert_eq!(results[0].record.id, second.id); - // assert!( - // results[0].distance <= results[1].distance, - // "results not ordered by distance: {:?}", - // results - // ); + let results = store.search(&query, Some(2)).await.unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].record.id, second.id); + assert!( + results[0].distance <= results[1].distance, + "results not ordered by distance: {:?}", + results + ); }); } @@ -1231,13 +1250,12 @@ mod tests { let store = ContextStore::open(&uri).await.unwrap(); // Verify we can list them back - let _results = store.list(None, None).await.unwrap(); - // TODO: MemWAL reads are not yet visible via standard scan. - // assert_eq!(results.len(), 2); + let results = store.list(None, None).await.unwrap(); + assert_eq!(results.len(), 2); - // let ids: Vec = results.iter().map(|r| r.id.clone()).collect(); - // assert!(ids.contains(&"r1".to_string())); - // assert!(ids.contains(&"r2".to_string())); + let ids: Vec = results.iter().map(|r| r.id.clone()).collect(); + assert!(ids.contains(&"r1".to_string())); + assert!(ids.contains(&"r2".to_string())); }); }