fix: search panic with prefilter and HNSW (#2584)

we call `search_in_partition` directly instead of `search` after recent changes, so move the waiting for prefilter ready to `search_in_partition` --------- Signed-off-by: BubbleCal <bubble-cal@outlook.com>
lancedb · Jul 12, 2024 · 0a9392e · 0a9392e
1 parent 5a8fb8c
commit 0a9392e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs
@@ -1731,6 +1731,9 @@ mod test {
     use datafusion::logical_expr::{col, lit};
     use half::f16;
     use lance_datagen::{array, gen, BatchCount, Dimension, RowCount};
+    use lance_index::vector::hnsw::builder::HnswBuildParams;
+    use lance_index::vector::ivf::IvfBuildParams;
+    use lance_index::vector::sq::builder::SQBuildParams;
     use lance_index::IndexType;
     use lance_io::object_store::ObjectStoreParams;
     use lance_testing::datagen::{BatchGenerator, IncrementingInt32, RandomVector};
@@ -2582,6 +2585,16 @@ mod test {
     async fn test_ann_prefilter(
         #[values(false, true)] use_legacy_format: bool,
         #[values(false, true)] stable_row_ids: bool,
+        #[values(
+            VectorIndexParams::ivf_pq(2, 8, 2, MetricType::L2, 2),
+            VectorIndexParams::with_ivf_hnsw_sq_params(
+                MetricType::L2,
+                IvfBuildParams::new(2),
+                HnswBuildParams::default(),
+                SQBuildParams::default()
+            )
+        )]
+        index_params: VectorIndexParams,
     ) {
         let test_dir = tempdir().unwrap();
         let test_uri = test_dir.path().to_str().unwrap();
@@ -2614,13 +2627,7 @@ mod test {
             .unwrap();
 
         dataset
-            .create_index(
-                &["vector"],
-                IndexType::Vector,
-                None,
-                &VectorIndexParams::ivf_pq(2, 8, 2, MetricType::L2, 2),
-                false,
-            )
+            .create_index(&["vector"], IndexType::Vector, None, &index_params, false)
             .await
             .unwrap();
 

diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs
@@ -352,7 +352,6 @@ impl<S: IvfSubIndex + fmt::Debug + 'static, Q: Quantization + fmt::Debug + 'stat
     for IVFIndex<S, Q>
 {
     async fn search(&self, query: &Query, pre_filter: Arc<dyn PreFilter>) -> Result<RecordBatch> {
-        pre_filter.wait_for_ready().await?;
         let mut query = query.clone();
         if self.distance_type == DistanceType::Cosine {
             let key = normalize_arrow(&query.key)?;
@@ -417,6 +416,7 @@ impl<S: IvfSubIndex + fmt::Debug + 'static, Q: Quantization + fmt::Debug + 'stat
         pre_filter: Arc<dyn PreFilter>,
     ) -> Result<RecordBatch> {
         let part_entry = self.load_partition(partition_id, true).await?;
+        pre_filter.wait_for_ready().await?;
 
         let query = self.preprocess_query(partition_id, query)?;
         let param = (&query).into();