mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
test: cover HNSW index optimization
Add a regression test for the public optimize-all path on a table with an IVF_HNSW_SQ index. The test creates 512 rows, runs the full optimize action, and verifies that the single retained index still covers all 512 rows with zero unindexed rows. Test-only change; no production behavior changes.
This commit is contained in:
@@ -440,8 +440,11 @@ mod tests {
|
||||
|
||||
use crate::connect;
|
||||
use crate::database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS;
|
||||
use crate::index::vector::IvfRqIndexBuilder;
|
||||
use crate::index::{Index, scalar::BTreeIndexBuilder};
|
||||
use crate::index::{
|
||||
Index,
|
||||
scalar::BTreeIndexBuilder,
|
||||
vector::{IvfHnswSqIndexBuilder, IvfRqIndexBuilder},
|
||||
};
|
||||
use crate::query::ExecutableQuery;
|
||||
use crate::table::{CompactionOptions, OptimizeAction, OptimizeStats};
|
||||
use futures::TryStreamExt;
|
||||
@@ -1178,6 +1181,60 @@ mod tests {
|
||||
assert_eq!(all_values, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_all_with_hnsw_index() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
|
||||
let dimension = 8;
|
||||
let item_field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(item_field.clone(), dimension),
|
||||
false,
|
||||
)]));
|
||||
|
||||
let make_batch = |offset: usize| {
|
||||
let values = Float32Array::from_iter_values(
|
||||
(offset * dimension as usize..(offset + 128) * dimension as usize)
|
||||
.map(|value| value as f32),
|
||||
);
|
||||
let vectors =
|
||||
FixedSizeListArray::try_new(item_field.clone(), dimension, Arc::new(values), None)
|
||||
.unwrap();
|
||||
RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap()
|
||||
};
|
||||
|
||||
let table = conn
|
||||
.create_table("test_hnsw_optimize", make_batch(0))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
for offset in [128, 256, 384] {
|
||||
table.add(make_batch(offset)).execute().await.unwrap();
|
||||
}
|
||||
|
||||
table
|
||||
.create_index(
|
||||
&["vector"],
|
||||
Index::IvfHnswSq(IvfHnswSqIndexBuilder::default()),
|
||||
)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let stats = table.optimize(OptimizeAction::All).await.unwrap();
|
||||
assert!(stats.compaction.unwrap().fragments_removed > 0);
|
||||
|
||||
let indices = table.list_indices().await.unwrap();
|
||||
assert_eq!(indices.len(), 1);
|
||||
assert_eq!(indices[0].index_type, crate::index::IndexType::IvfHnswSq);
|
||||
|
||||
let index_stats = table.index_stats(&indices[0].name).await.unwrap().unwrap();
|
||||
assert_eq!(index_stats.num_indexed_rows, 512);
|
||||
assert_eq!(index_stats.num_unindexed_rows, 0);
|
||||
assert_eq!(table.count_rows(None).await.unwrap(), 512);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_default_action() {
|
||||
// Verify that default action is All
|
||||
|
||||
Reference in New Issue
Block a user