From f52f07539364d5b62bd27c1ed1e469d80602e8b4 Mon Sep 17 00:00:00 2001 From: Geser Dugarov Date: Tue, 7 Jul 2026 11:11:10 +0700 Subject: [PATCH 1/2] fix(index): keep HNSW partition scratch dir alive until tasks finish --- rust/lance/src/index/vector/ivf/io.rs | 584 +++++++++++++++++++------- 1 file changed, 442 insertions(+), 142 deletions(-) diff --git a/rust/lance/src/index/vector/ivf/io.rs b/rust/lance/src/index/vector/ivf/io.rs index 612dd162281..d2aec717b69 100644 --- a/rust/lance/src/index/vector/ivf/io.rs +++ b/rust/lance/src/index/vector/ivf/io.rs @@ -44,6 +44,7 @@ use lance_table::format::SelfDescribingFileReader; use lance_table::io::manifest::ManifestDescribing; use object_store::path::Path; use tokio::sync::Semaphore; +use tokio::task::JoinHandle; use crate::Result; @@ -278,164 +279,222 @@ pub(super) async fn write_hnsw_quantization_index_partitions( } let object_store = ObjectStore::local(); - let mut part_files = Vec::with_capacity(ivf.num_partitions()); - let mut aux_part_files = Vec::with_capacity(ivf.num_partitions()); - let tmp_part_dir = Path::from_filesystem_path(TempStdDir::default())?; - let mut tasks = Vec::with_capacity(ivf.num_partitions()); - let sem = Arc::new(Semaphore::new(*HNSW_PARTITIONS_BUILD_PARALLEL)); - for part_id in 0..ivf.num_partitions() { - part_files.push(tmp_part_dir.clone().join(format!("hnsw_part_{}", part_id))); - aux_part_files.push( - tmp_part_dir - .clone() - .join(format!("hnsw_part_aux_{}", part_id)), - ); + // Partitions are staged in this scratch dir, then merged into the final index. + // Share the guard with every task via `Arc` so its `Drop` removes the dir only + // after the last task finishes -- never while one is still writing. + let tmp_part_dir_guard = Arc::new(TempStdDir::default()); + let tmp_part_dir = Path::from_filesystem_path(&**tmp_part_dir_guard)?; + + // `Option` per handle so the consume loop can `take()` each one, leaving the + // not-yet-consumed handles for the error-path drain. + let mut tasks: Vec>>> = + Vec::with_capacity(ivf.num_partitions()); + + let build_result: Result<(Vec, IvfModel)> = async { + let mut part_files = Vec::with_capacity(ivf.num_partitions()); + let mut aux_part_files = Vec::with_capacity(ivf.num_partitions()); + let sem = Arc::new(Semaphore::new(*HNSW_PARTITIONS_BUILD_PARALLEL)); + for part_id in 0..ivf.num_partitions() { + part_files.push(tmp_part_dir.clone().join(format!("hnsw_part_{}", part_id))); + aux_part_files.push( + tmp_part_dir + .clone() + .join(format!("hnsw_part_aux_{}", part_id)), + ); - let mut code_array: Vec> = vec![]; - let mut row_id_array: Vec> = vec![]; + let mut code_array: Vec> = vec![]; + let mut row_id_array: Vec> = vec![]; - // We don't transform vectors to SQ codes while shuffling, - // so we won't merge SQ codes from the stream. + // We don't transform vectors to SQ codes while shuffling, + // so we won't merge SQ codes from the stream. - if let Some(&previous_indices) = existing_indices.as_ref() { - for &idx in previous_indices.iter() { - let sub_index = idx - .load_partition(part_id, true, &NoOpMetricsCollector) - .await?; - let row_ids = Arc::new(UInt64Array::from_iter_values(sub_index.row_ids().cloned())); - row_id_array.push(row_ids); + if let Some(&previous_indices) = existing_indices.as_ref() { + for &idx in previous_indices.iter() { + let sub_index = idx + .load_partition(part_id, true, &NoOpMetricsCollector) + .await?; + let row_ids = + Arc::new(UInt64Array::from_iter_values(sub_index.row_ids().cloned())); + row_id_array.push(row_ids); + } } - } - let code_column = match &quantizer { - Quantizer::Product(pq) => Some(pq.column()), - _ => None, - }; - merge_streams( - &mut streams_heap, - &mut new_streams, - part_id as u32, - code_column, - &mut code_array, - &mut row_id_array, - ) - .await?; - - if row_id_array.is_empty() { - tasks.push(tokio::spawn(async { Ok(0) })); - continue; - } - - let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); - let part_writer = PreviousFileWriter::::try_new( - &object_store, - part_file, - Schema::try_from(writer.schema())?, - &Default::default(), - ) - .await?; + let code_column = match &quantizer { + Quantizer::Product(pq) => Some(pq.column()), + _ => None, + }; + merge_streams( + &mut streams_heap, + &mut new_streams, + part_id as u32, + code_column, + &mut code_array, + &mut row_id_array, + ) + .await?; - let aux_part_writer = match auxiliary_writer.as_ref() { - Some(writer) => Some( - PreviousFileWriter::::try_new( - &object_store, - aux_part_file, - Schema::try_from(writer.schema())?, - &Default::default(), - ) - .await?, - ), - None => None, - }; + if row_id_array.is_empty() { + tasks.push(Some(tokio::spawn(async { Ok(0) }))); + continue; + } - let dataset = dataset.clone(); - let column = column.to_owned(); - let hnsw_params = hnsw_params.clone(); - let quantizer = quantizer.clone(); - let sem = sem.clone(); - tasks.push(tokio::spawn(async move { - let _permit = sem.acquire().await.expect("semaphore error"); - - log::debug!("Building HNSW partition {}", part_id); - let result = build_hnsw_quantization_partition( - dataset, - &column, - distance_type, - hnsw_params, - part_writer, - aux_part_writer, - quantizer, - row_id_array, - code_array, + let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); + let part_writer = PreviousFileWriter::::try_new( + &object_store, + part_file, + Schema::try_from(writer.schema())?, + &Default::default(), ) - .await; - log::debug!("Finished building HNSW partition {}", part_id); - result - })); - } - - let mut aux_ivf = IvfModel::empty(); - let mut hnsw_metadata = Vec::with_capacity(ivf.num_partitions()); - for (part_id, task) in tasks.into_iter().enumerate() { - let offset = writer.len(); - let num_rows = task.await??; + .await?; - if num_rows == 0 { - ivf.add_partition(0); - aux_ivf.add_partition(0); - hnsw_metadata.push(HnswMetadata::default()); - continue; + let aux_part_writer = match auxiliary_writer.as_ref() { + Some(writer) => Some( + PreviousFileWriter::::try_new( + &object_store, + aux_part_file, + Schema::try_from(writer.schema())?, + &Default::default(), + ) + .await?, + ), + None => None, + }; + + let dataset = dataset.clone(); + let column = column.to_owned(); + let hnsw_params = hnsw_params.clone(); + let quantizer = quantizer.clone(); + let sem = sem.clone(); + let tmp_part_dir_guard = tmp_part_dir_guard.clone(); + tasks.push(Some(tokio::spawn(async move { + // Hold a guard clone so the scratch dir stays alive while this task writes. + let _tmp_part_dir_guard = tmp_part_dir_guard; + let _permit = sem.acquire().await.expect("semaphore error"); + + log::debug!("Building HNSW partition {}", part_id); + let result = build_hnsw_quantization_partition( + dataset, + &column, + distance_type, + hnsw_params, + part_writer, + aux_part_writer, + quantizer, + row_id_array, + code_array, + ) + .await; + log::debug!("Finished building HNSW partition {}", part_id); + result + }))); } - let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); - let part_reader = - PreviousFileReader::try_new_self_described(&object_store, part_file, None).await?; + let mut aux_ivf = IvfModel::empty(); + let mut hnsw_metadata = Vec::with_capacity(ivf.num_partitions()); + for (part_id, task) in tasks.iter_mut().enumerate() { + let task = task + .take() + .expect("each partition task is consumed exactly once"); + let offset = writer.len(); + let num_rows = task.await??; + + if num_rows == 0 { + ivf.add_partition(0); + aux_ivf.add_partition(0); + hnsw_metadata.push(HnswMetadata::default()); + continue; + } - let batches = futures::stream::iter(0..part_reader.num_batches()) - .map(|batch_id| { - part_reader.read_batch( - batch_id as i32, - ReadBatchParams::RangeFull, - part_reader.schema(), - ) - }) - .buffered(object_store.io_parallelism()) - .try_collect::>() - .await?; - writer.write(&batches).await?; - - ivf.add_partition((writer.len() - offset) as u32); - hnsw_metadata.push(serde_json::from_str( - part_reader.schema().metadata[HNSW::metadata_key()].as_str(), - )?); - std::mem::drop(part_reader); - object_store.delete(part_file).await?; - - if let Some(aux_writer) = auxiliary_writer.as_mut() { - let aux_part_reader = - PreviousFileReader::try_new_self_described(&object_store, aux_part_file, None) - .await?; + let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); + let part_reader = + PreviousFileReader::try_new_self_described(&object_store, part_file, None).await?; - let batches = futures::stream::iter(0..aux_part_reader.num_batches()) + let batches = futures::stream::iter(0..part_reader.num_batches()) .map(|batch_id| { - aux_part_reader.read_batch( + part_reader.read_batch( batch_id as i32, ReadBatchParams::RangeFull, - aux_part_reader.schema(), + part_reader.schema(), ) }) .buffered(object_store.io_parallelism()) .try_collect::>() .await?; - std::mem::drop(aux_part_reader); - object_store.delete(aux_part_file).await?; + writer.write(&batches).await?; + + ivf.add_partition((writer.len() - offset) as u32); + hnsw_metadata.push(serde_json::from_str( + part_reader.schema().metadata[HNSW::metadata_key()].as_str(), + )?); + std::mem::drop(part_reader); + object_store.delete(part_file).await?; + + if let Some(aux_writer) = auxiliary_writer.as_mut() { + let aux_part_reader = + PreviousFileReader::try_new_self_described(&object_store, aux_part_file, None) + .await?; + + let batches = futures::stream::iter(0..aux_part_reader.num_batches()) + .map(|batch_id| { + aux_part_reader.read_batch( + batch_id as i32, + ReadBatchParams::RangeFull, + aux_part_reader.schema(), + ) + }) + .buffered(object_store.io_parallelism()) + .try_collect::>() + .await?; + std::mem::drop(aux_part_reader); + object_store.delete(aux_part_file).await?; - aux_writer.write(&batches).await?; - aux_ivf.add_partition(num_rows as u32); + aux_writer.write(&batches).await?; + aux_ivf.add_partition(num_rows as u32); + } + } + + Ok((hnsw_metadata, aux_ivf)) + } + .await; + + // On error, abort and await the partition builds we never consumed so none + // keep running in the background; see `drain_partition_tasks`. + if build_result.is_err() { + for err in drain_partition_tasks(&mut tasks).await { + log::warn!( + "HNSW partition build task failed while draining after an earlier error: {err}" + ); } } - Ok((hnsw_metadata, aux_ivf)) + build_result +} + +/// Abort and await every still-outstanding partition-build task, returning the +/// non-cancellation errors they surfaced. +/// +/// A dropped [`JoinHandle`] detaches its task, so each handle is aborted and +/// awaited: awaiting resolves only once the task has actually stopped and is no +/// longer writing into the scratch dir. Cancellation errors are the expected +/// result of the abort and dropped; real failures and panics are returned so the +/// caller can surface them. +async fn drain_partition_tasks(tasks: &mut [Option>>]) -> Vec { + let mut errors = Vec::new(); + for task in tasks.iter_mut() { + let Some(handle) = task.take() else { + continue; + }; + handle.abort(); + match handle.await { + Ok(Ok(_)) => {} + Ok(Err(e)) => errors.push(e), + Err(join_err) if join_err.is_cancelled() => {} + Err(join_err) => errors.push(Error::io(format!( + "HNSW partition build task panicked: {join_err}" + ))), + } + } + errors } #[allow(clippy::too_many_arguments)] @@ -473,24 +532,22 @@ async fn build_hnsw_quantization_partition( let build_hnsw = build_and_write_hnsw(vectors.clone(), (*hnsw_params).clone(), metric_type, writer); + // Build PQ storage as a child future, joined below: it writes `aux_writer`'s + // file into the scratch dir, so it is cancelled together with this task when the + // error-path drain aborts it, and the join surfaces its errors. let build_store = match quantizer { Quantizer::Flat(_) => { return Err(Error::index( "Flat quantizer is not supported for IVF_HNSW".to_string(), )); } - Quantizer::Product(pq) => tokio::spawn(build_and_write_pq_storage( - metric_type, - row_ids, - code_array, - pq, - aux_writer.unwrap(), - )), - + Quantizer::Product(pq) => { + build_and_write_pq_storage(metric_type, row_ids, code_array, pq, aux_writer.unwrap()) + } _ => unreachable!("IVF_HNSW_SQ has been moved to v2 index builder"), }; - let index_rows = futures::join!(build_hnsw, build_store).0?; + let (index_rows, ()) = futures::try_join!(build_hnsw, build_store)?; assert!( index_rows >= num_rows, "index rows {} must be greater than or equal to num rows {}", @@ -530,6 +587,8 @@ async fn build_and_write_pq_storage( mod tests { use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use crate::Dataset; use crate::index::vector::ivf::v2; use crate::index::{DatasetIndexExt, DatasetIndexInternalExt, vector::VectorIndexParams}; @@ -589,4 +648,245 @@ mod tests { //let indices = /ds. } + + /// The scratch dir must outlive every partition task: dropping the caller-side + /// guard while tasks are still in flight must not remove it, because each task + /// still holds an `Arc` clone of the guard. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_scratch_dir_outlives_partition_tasks() { + let tmp_part_dir_guard = Arc::new(TempStdDir::default()); + let scratch_path = tmp_part_dir_guard.to_path_buf(); + + // Park each task until the caller-side guard is dropped, so the dir's + // survival is attributable solely to the clones the tasks still hold. + let running = Arc::new(AtomicUsize::new(0)); + let released = Arc::new(AtomicBool::new(false)); + + const NUM_TASKS: usize = 3; + let mut tasks = Vec::with_capacity(NUM_TASKS); + for _ in 0..NUM_TASKS { + let task_guard = tmp_part_dir_guard.clone(); + let running = running.clone(); + let released = released.clone(); + tasks.push(tokio::spawn(async move { + running.fetch_add(1, Ordering::SeqCst); + while !released.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + // Still holding `task_guard`, so the directory must be live. + task_guard.exists() + })); + } + + // Drop the caller-side guard only once every task is parked holding its clone. + while running.load(Ordering::SeqCst) < NUM_TASKS { + tokio::task::yield_now().await; + } + drop(tmp_part_dir_guard); + assert!( + scratch_path.exists(), + "scratch dir removed while partition tasks still held the guard" + ); + + released.store(true, Ordering::SeqCst); + for task in tasks { + assert!( + task.await.unwrap(), + "a partition task observed its scratch dir already removed" + ); + } + assert!( + !scratch_path.exists(), + "scratch dir was not removed after the last task's guard clone dropped" + ); + } + + /// The drain step must outlive every spawned task: it aborts and awaits each + /// one so that, once it returns, no task is still running against the scratch + /// directory. It must also surface late failures rather than swallow them. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_drain_partition_tasks_waits_and_reports_errors() { + // Scratch dir guard analogous to the one held by + // `write_hnsw_quantization_index_partitions`; tasks read from it, and it + // is dropped only after the drain completes. + let scratch_guard = TempStdDir::default(); + let scratch_path = scratch_guard.to_path_buf(); + + // Count of live task futures. `LiveGuard` decrements on both completion and + // cancellation, so a zero count after the drain proves every task terminated + // rather than being detached. + let live = Arc::new(AtomicUsize::new(0)); + let saw_missing_dir = Arc::new(AtomicBool::new(false)); + + struct LiveGuard(Arc); + impl Drop for LiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + + const NUM_SLOW: usize = 3; + let mut tasks: Vec>>> = Vec::new(); + + // Tasks that never resolve on their own: the drain's abort is the only thing + // that can stop them, which is exactly what this test exercises. + for _ in 0..NUM_SLOW { + let live = live.clone(); + let saw_missing_dir = saw_missing_dir.clone(); + let scratch_path = scratch_path.clone(); + tasks.push(Some(tokio::spawn(async move { + live.fetch_add(1, Ordering::SeqCst); + let _guard = LiveGuard(live.clone()); + if !scratch_path.exists() { + saw_missing_dir.store(true, Ordering::SeqCst); + } + futures::future::pending::<()>().await; + Ok(0) + }))); + } + + // A task that fails; its error must be returned, not silently dropped. + tasks.push(Some(tokio::spawn(async move { + Err(Error::io("late partition failure".to_string())) + }))); + + // Ensure all slow tasks are actually running before draining, so the + // drain has to await their cancellation rather than aborting them before + // they ever start. + while live.load(Ordering::SeqCst) < NUM_SLOW { + tokio::task::yield_now().await; + } + + let errors = drain_partition_tasks(&mut tasks).await; + + assert!(tasks.iter().all(Option::is_none), "handles left undrained"); + assert_eq!( + live.load(Ordering::SeqCst), + 0, + "a task was still running after the drain returned" + ); + assert!( + !saw_missing_dir.load(Ordering::SeqCst), + "scratch dir was removed while a task was still running" + ); + assert_eq!(errors.len(), 1, "expected exactly the one late failure"); + assert!( + errors[0].to_string().contains("late partition failure"), + "late failure was not surfaced: {}", + errors[0] + ); + + // The guard, not the drain, removes the scratch dir. + assert!(scratch_path.exists()); + drop(scratch_guard); + assert!(!scratch_path.exists()); + } + + /// `write_hnsw_quantization_index_partitions` stages each partition in a scratch + /// directory owned by a [`TempStdDir`] guard, which must remove it once the build + /// finishes so the OS temp dir does not grow without bound across legacy + /// IVF_HNSW_* builds. + /// + /// The OS temp dir is process-global, so an in-process check can't attribute a + /// leak to our own build. Instead we run the build in a child process with + /// `TMPDIR` pointed at an isolated dir we own, then assert nothing survives. + #[test] + fn test_hnsw_pq_scratch_dir_is_not_leaked() { + use std::path::PathBuf; + + // Isolated temp root for the child. Owned here so it -- and anything the + // child leaks into it -- is removed when this guard drops at test end. + let isolated_root = TempStdDir::default(); + + let child_test = "index::vector::ivf::io::tests::build_legacy_hnsw_pq_in_child_process"; + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([child_test, "--exact", "--ignored", "--nocapture"]) + .env("TMPDIR", isolated_root.as_ref()) + .env("LANCE_HNSW_LEAK_TEST_ROOT", isolated_root.as_ref()) + .output() + .expect("failed to spawn child test process"); + assert!( + output.status.success(), + "child build process failed:\n--- stdout ---\n{}\n--- stderr ---\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + + // Each build stages its partitions in one `.tmp*` dir directly under TMPDIR. + // Every guard removes its dir when it drops, so none should survive. + let leaked: Vec = std::fs::read_dir(&isolated_root) + .expect("read isolated temp root") + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.is_dir() + && path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.starts_with(".tmp")) + }) + .collect(); + + assert!( + leaked.is_empty(), + "legacy IVF_HNSW_PQ build leaked {} scratch director{} under the temp dir; \ + the TempStdDir guard should remove each one when it drops: {:?}", + leaked.len(), + if leaked.len() == 1 { "y" } else { "ies" }, + leaked, + ); + } + + /// Child half of [`test_hnsw_pq_scratch_dir_is_not_leaked`]. Ignored so it only + /// runs when the parent spawns it with `TMPDIR` and `LANCE_HNSW_LEAK_TEST_ROOT` + /// pointed at an isolated dir. Builds a few legacy IVF_HNSW_PQ indices; the + /// parent does the leak detection. + #[tokio::test] + #[ignore = "spawned as a child process by test_hnsw_pq_scratch_dir_is_not_leaked"] + async fn build_legacy_hnsw_pq_in_child_process() { + use lance_index::vector::ivf::IvfBuildParams; + use lance_index::vector::pq::PQBuildParams; + + // Only do work when spawned by the parent; a bare `--ignored` run leaves + // the variable unset, so no-op rather than fail. + let Ok(root) = std::env::var("LANCE_HNSW_LEAK_TEST_ROOT") else { + return; + }; + + const DIM: usize = 32; + const ROWS: usize = 1024; + const NLIST: usize = 4; + const NUM_BUILDS: usize = 3; + + // Keep the dataset out of the temp dir's `.tmp*` namespace so the parent + // never confuses it with a leaked scratch directory. + let dataset_uri = format!("{root}/dataset"); + let values = generate_random_array(ROWS * DIM); + let fsl = Arc::new(FixedSizeListArray::try_new_from_values(values, DIM as i32).unwrap()); + let schema = Arc::new(Schema::new(vec![Field::new( + "vector", + fsl.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![fsl]).unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let ds = Dataset::write(batches, &dataset_uri, Default::default()) + .await + .unwrap(); + + for _ in 0..NUM_BUILDS { + crate::index::vector::ivf::build_ivf_hnsw_pq_index( + &ds, + "vector", + "idx", + uuid::Uuid::new_v4(), + MetricType::L2, + &IvfBuildParams::new(NLIST), + &HnswBuildParams::default(), + &PQBuildParams::new(4, 8), + ) + .await + .unwrap(); + } + } } From c76e78c942797de2b482f2098d8f5bcfd6544737 Mon Sep 17 00:00:00 2001 From: Geser Dugarov Date: Thu, 9 Jul 2026 16:19:33 +0700 Subject: [PATCH 2/2] fix(index): abort HNSW partition tasks before draining --- rust/lance/src/index/vector/ivf/io.rs | 63 +++++++++++++++++++-------- 1 file changed, 44 insertions(+), 19 deletions(-) diff --git a/rust/lance/src/index/vector/ivf/io.rs b/rust/lance/src/index/vector/ivf/io.rs index d2aec717b69..f8f713a07af 100644 --- a/rust/lance/src/index/vector/ivf/io.rs +++ b/rust/lance/src/index/vector/ivf/io.rs @@ -369,7 +369,11 @@ pub(super) async fn write_hnsw_quantization_index_partitions( tasks.push(Some(tokio::spawn(async move { // Hold a guard clone so the scratch dir stays alive while this task writes. let _tmp_part_dir_guard = tmp_part_dir_guard; - let _permit = sem.acquire().await.expect("semaphore error"); + let _permit = sem.acquire().await.map_err(|err| { + Error::io(format!( + "failed to acquire HNSW partition build permit: {err}" + )) + })?; log::debug!("Building HNSW partition {}", part_id); let result = build_hnsw_quantization_partition( @@ -473,18 +477,26 @@ pub(super) async fn write_hnsw_quantization_index_partitions( /// Abort and await every still-outstanding partition-build task, returning the /// non-cancellation errors they surfaced. /// -/// A dropped [`JoinHandle`] detaches its task, so each handle is aborted and -/// awaited: awaiting resolves only once the task has actually stopped and is no -/// longer writing into the scratch dir. Cancellation errors are the expected -/// result of the abort and dropped; real failures and panics are returned so the -/// caller can surface them. +/// A dropped [`JoinHandle`] detaches its task, so every handle is aborted up front +/// before any is awaited: otherwise a task slow to observe its own cancellation +/// would keep running -- and keep writing into the scratch dir -- while an earlier +/// handle is still being awaited. Awaiting then resolves only once each task has +/// actually stopped. Cancellation errors are the expected result of the abort and +/// dropped; failures and panics from tasks that had already finished before the +/// abort are returned so the caller can surface them (a task still in flight is +/// cancelled, so this best-effort drain only reports errors already produced). async fn drain_partition_tasks(tasks: &mut [Option>>]) -> Vec { - let mut errors = Vec::new(); + for task in tasks.iter() { + if let Some(handle) = task.as_ref() { + handle.abort(); + } + } + + let mut errors = Vec::with_capacity(tasks.len()); for task in tasks.iter_mut() { let Some(handle) = task.take() else { continue; }; - handle.abort(); match handle.await { Ok(Ok(_)) => {} Ok(Err(e)) => errors.push(e), @@ -542,9 +554,16 @@ async fn build_hnsw_quantization_partition( )); } Quantizer::Product(pq) => { - build_and_write_pq_storage(metric_type, row_ids, code_array, pq, aux_writer.unwrap()) + let aux_writer = aux_writer.ok_or_else(|| { + Error::index("IVF_HNSW_PQ requires an auxiliary writer for PQ storage".to_string()) + })?; + build_and_write_pq_storage(metric_type, row_ids, code_array, pq, aux_writer) + } + _ => { + return Err(Error::index( + "IVF_HNSW_SQ is not supported in the legacy HNSW partition writer".to_string(), + )); } - _ => unreachable!("IVF_HNSW_SQ has been moved to v2 index builder"), }; let (index_rows, ()) = futures::try_join!(build_hnsw, build_store)?; @@ -587,6 +606,7 @@ async fn build_and_write_pq_storage( mod tests { use super::*; + use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use crate::Dataset; @@ -597,6 +617,8 @@ mod tests { use lance_core::utils::tempfile::TempStrDir; use lance_index::IndexType; use lance_index::metrics::NoOpMetricsCollector; + use lance_index::vector::ivf::IvfBuildParams; + use lance_index::vector::pq::PQBuildParams; use lance_testing::datagen::generate_random_array; #[tokio::test] @@ -726,7 +748,7 @@ mod tests { } const NUM_SLOW: usize = 3; - let mut tasks: Vec>>> = Vec::new(); + let mut tasks: Vec>>> = Vec::with_capacity(NUM_SLOW + 1); // Tasks that never resolve on their own: the drain's abort is the only thing // that can stop them, which is exactly what this test exercises. @@ -746,9 +768,17 @@ mod tests { } // A task that fails; its error must be returned, not silently dropped. - tasks.push(Some(tokio::spawn(async move { - Err(Error::io("late partition failure".to_string())) - }))); + // The drain aborts every handle up front, and an abort only preserves a + // task's output if the task has already finished -- an in-flight task is + // cancelled and its error lost. So wait until this task has actually + // completed before handing it to the drain; otherwise whether its error + // surfaces would depend on the scheduler and the test would be flaky. + let failing = + tokio::spawn(async move { Err(Error::io("late partition failure".to_string())) }); + while !failing.is_finished() { + tokio::task::yield_now().await; + } + tasks.push(Some(failing)); // Ensure all slow tasks are actually running before draining, so the // drain has to await their cancellation rather than aborting them before @@ -792,8 +822,6 @@ mod tests { /// `TMPDIR` pointed at an isolated dir we own, then assert nothing survives. #[test] fn test_hnsw_pq_scratch_dir_is_not_leaked() { - use std::path::PathBuf; - // Isolated temp root for the child. Owned here so it -- and anything the // child leaks into it -- is removed when this guard drops at test end. let isolated_root = TempStdDir::default(); @@ -844,9 +872,6 @@ mod tests { #[tokio::test] #[ignore = "spawned as a child process by test_hnsw_pq_scratch_dir_is_not_leaked"] async fn build_legacy_hnsw_pq_in_child_process() { - use lance_index::vector::ivf::IvfBuildParams; - use lance_index::vector::pq::PQBuildParams; - // Only do work when spawned by the parent; a bare `--ignored` run leaves // the variable unset, so no-op rather than fail. let Ok(root) = std::env::var("LANCE_HNSW_LEAK_TEST_ROOT") else {