From 8ad23d441a4166214f08d919753ca7fff646c732 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 20 Oct 2026 11:14:00 +0800 Subject: [PATCH 1/2] fix: resolve vector prewarm name collisions from segment metadata --- src/index.rs | 45 +++++++++++++++++++++++++-- tests/c_api_test.rs | 74 ++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 112 insertions(+), 7 deletions(-) diff --git a/src/index.rs b/src/index.rs index c3add47..4b0f8b5 100644 --- a/src/index.rs +++ b/src/index.rs @@ -10,9 +10,10 @@ use std::ffi::{CString, c_char}; use std::sync::Arc; use arrow_array::{Array, ArrayRef, FixedSizeListArray}; -use lance::index::DatasetIndexExt; +use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; use lance_core::Result; use lance_index::IndexType; +use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; use crate::dataset::LanceDataset; @@ -297,10 +298,50 @@ unsafe fn prewarm_index_inner( // Keep one snapshot and its shared session alive without holding the // dataset lock across potentially long-running index IO. let snapshot = unsafe { &*dataset }.snapshot(); - block_on(snapshot.prewarm_index(index_name))?; + block_on(prewarm_index_by_name(&snapshot, index_name))?; Ok(0) } +async fn prewarm_index_by_name(dataset: &lance::Dataset, index_name: &str) -> Result<()> { + let Some(named_field) = dataset.schema().field(index_name) else { + return dataset.prewarm_index(index_name).await; + }; + let indices = dataset.load_indices_by_name(index_name).await?; + if indices.is_empty() { + return dataset.prewarm_index(index_name).await; + } + + let mut sdk_segments = Vec::new(); + for metadata in indices { + let Some(&field_id) = metadata.fields.first() else { + sdk_segments.push(metadata.uuid); + continue; + }; + if field_id == named_field.id { + sdk_segments.push(metadata.uuid); + continue; + } + // The pinned SDK interprets an index name as a schema column first. + // Open colliding vector segments by their indexed field, even on a cold cache. + let field_path = dataset.schema().field_path(field_id)?; + let index = dataset + .open_generic_index(&field_path, &metadata.uuid, &NoOpMetricsCollector) + .await?; + if index.index_type().is_vector() { + index.prewarm().await?; + } else { + sdk_segments.push(metadata.uuid); + } + } + // Keep scalar/FTS validation, including aggregate FTS residency checks, in the SDK. + if !sdk_segments.is_empty() { + dataset + .prewarm_index_segments(index_name, &sdk_segments) + .await?; + } + Ok(()) +} + /// Drop an index by name. #[unsafe(no_mangle)] pub unsafe extern "C" fn lance_dataset_drop_index( diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 341e553..2cd5d2b 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -18129,11 +18129,55 @@ fn test_prewarm_index_invalid_arguments_and_missing_index() { .contains("missing_idx") ); lance_free_string(message); + assert_eq!(lance_dataset_prewarm_index(ds, c"id".as_ptr()), -1); + assert_eq!(lance_last_error_code(), LanceErrorCode::NotFound); assert_eq!(lance_dataset_count_rows(ds), 5); lance_dataset_close(ds); } } +#[test] +fn test_prewarm_index_scalar_name_collides_with_other_column() { + for index_type in [LanceScalarIndexType::BTree, LanceScalarIndexType::Inverted] { + let (_tmp, uri) = create_test_dataset(); + unsafe { + let ds = lance_dataset_open(c_str(&uri).as_ptr(), ptr::null(), 0); + assert!(!ds.is_null()); + let params = c_str(r#"{"base_tokenizer":"simple","language":"English"}"#); + assert_eq!( + lance_dataset_create_scalar_index( + ds, + c"name".as_ptr(), + c"id".as_ptr(), + index_type as i32, + if index_type == LanceScalarIndexType::Inverted { + params.as_ptr() + } else { + ptr::null() + }, + false, + ), + 0, + "{}", + take_last_error_message() + ); + let version = lance_dataset_version(ds); + lance_dataset_close(ds); + // Scalar and FTS prewarm must retain their SDK validation on a fresh session too. + let cold = lance_dataset_open(c_str(&uri).as_ptr(), ptr::null(), version); + assert!(!cold.is_null()); + assert_eq!( + lance_dataset_prewarm_index(cold, c"id".as_ptr()), + 0, + "{}", + take_last_error_message() + ); + assert_eq!(lance_dataset_version(cold), version); + lance_dataset_close(cold); + } + } +} + #[test] fn test_prewarm_index_scalar_segments_and_snapshot() { let (_tmp, uri, _) = create_scalar_segment_fixture(lance_index::IndexType::BTree, false); @@ -18162,6 +18206,21 @@ fn test_prewarm_index_scalar_segments_and_snapshot() { #[test] fn test_prewarm_index_vector_segments_reuse_shared_session() { + check_prewarm_index_vector_segments_reuse_shared_session(c"vector_idx"); +} + +#[test] +fn test_prewarm_index_vector_name_matches_indexed_column() { + check_prewarm_index_vector_segments_reuse_shared_session(c"embedding"); +} + +#[test] +fn test_prewarm_index_vector_name_collides_with_scalar_column() { + // A fresh session must resolve the vector field from metadata, not the scalar column named id. + check_prewarm_index_vector_segments_reuse_shared_session(c"id"); +} + +fn check_prewarm_index_vector_segments_reuse_shared_session(index_name: &std::ffi::CStr) { use lance::index::{DatasetIndexExt, vector::VectorIndexParams}; let (_tmp, uri) = create_multi_fragment_vector_dataset(2, 64, 8, false); lance_c::runtime::block_on(async { @@ -18171,14 +18230,14 @@ fn test_prewarm_index_vector_segments_reuse_shared_session() { for fragment in ds.get_fragments() { segments.push( ds.create_index_builder(&["embedding"], lance_index::IndexType::Vector, ¶ms) - .name("vector_idx".into()) + .name(index_name.to_str().unwrap().into()) .fragments(vec![fragment.id() as u32]) .execute_uncommitted() .await .unwrap(), ); } - ds.commit_existing_index_segments("vector_idx", "embedding", segments) + ds.commit_existing_index_segments(index_name.to_str().unwrap(), "embedding", segments) .await .unwrap(); }); @@ -18240,11 +18299,16 @@ fn test_prewarm_index_vector_segments_reuse_shared_session() { let ds = lance_dataset_open_with_session(c_str(&uri).as_ptr(), ptr::null(), 0, session); assert!(!ds.is_null()); assert_eq!( - lance_dataset_index_segment_count(ds, c"vector_idx".as_ptr()), + lance_dataset_index_segment_count(ds, index_name.as_ptr()), 2 ); let version = lance_dataset_version(ds); - assert_eq!(lance_dataset_prewarm_index(ds, c"vector_idx".as_ptr()), 0); + assert_eq!( + lance_dataset_prewarm_index(ds, index_name.as_ptr()), + 0, + "{}", + take_last_error_message() + ); assert_eq!(lance_dataset_version(ds), version); let mut warmed = LanceSessionCacheStats::default(); assert_eq!(lance_session_get_cache_stats(session, &mut warmed), 0); @@ -18274,7 +18338,7 @@ fn test_prewarm_index_vector_segments_reuse_shared_session() { // even after its C session handle is released. lance_session_close(session); assert_eq!( - lance_dataset_prewarm_index(reopened, c"vector_idx".as_ptr()), + lance_dataset_prewarm_index(reopened, index_name.as_ptr()), 0 ); assert_eq!(lance_dataset_version(reopened), version); From 91f80585a381c59cbe010217800b80d16bf963ca Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 20 Oct 2026 14:21:57 +0800 Subject: [PATCH 2/2] fix: preserve concurrent prewarm for colliding index segments Run each segment open and prewarm operation together under try_join_all, matching the SDK concurrency and error handling. Keep scalar and unresolved segments on the SDK validation path. Add deterministic gate-based tests for concurrent opens and prewarms, error cancellation, and SDK fallback routing. --- src/index.rs | 73 ++++++++++++++------- src/index_tests.rs | 154 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 205 insertions(+), 22 deletions(-) create mode 100644 src/index_tests.rs diff --git a/src/index.rs b/src/index.rs index 4b0f8b5..e53ffe9 100644 --- a/src/index.rs +++ b/src/index.rs @@ -7,6 +7,7 @@ //! scanners that already cloned the inner Arc keep their snapshot view. use std::ffi::{CString, c_char}; +use std::future::Future; use std::sync::Arc; use arrow_array::{Array, ArrayRef, FixedSizeListArray}; @@ -15,6 +16,8 @@ use lance_core::Result; use lance_index::IndexType; use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; +use lance_table::format::IndexMetadata; +use uuid::Uuid; use crate::dataset::LanceDataset; use crate::error::ffi_try; @@ -311,28 +314,22 @@ async fn prewarm_index_by_name(dataset: &lance::Dataset, index_name: &str) -> Re return dataset.prewarm_index(index_name).await; } - let mut sdk_segments = Vec::new(); - for metadata in indices { - let Some(&field_id) = metadata.fields.first() else { - sdk_segments.push(metadata.uuid); - continue; - }; - if field_id == named_field.id { - sdk_segments.push(metadata.uuid); - continue; - } - // The pinned SDK interprets an index name as a schema column first. - // Open colliding vector segments by their indexed field, even on a cold cache. - let field_path = dataset.schema().field_path(field_id)?; - let index = dataset - .open_generic_index(&field_path, &metadata.uuid, &NoOpMetricsCollector) - .await?; - if index.index_type().is_vector() { - index.prewarm().await?; - } else { - sdk_segments.push(metadata.uuid); - } - } + let sdk_segments = + prewarm_colliding_segments(indices, named_field.id, |field_id, uuid| async move { + // The pinned SDK interprets an index name as a schema column first. + // Open colliding vector segments by their indexed field, even on a cold cache. + let field_path = dataset.schema().field_path(field_id)?; + let index = dataset + .open_generic_index(&field_path, &uuid, &NoOpMetricsCollector) + .await?; + if index.index_type().is_vector() { + index.prewarm().await?; + Ok(true) + } else { + Ok(false) + } + }) + .await?; // Keep scalar/FTS validation, including aggregate FTS residency checks, in the SDK. if !sdk_segments.is_empty() { dataset @@ -342,6 +339,38 @@ async fn prewarm_index_by_name(dataset: &lance::Dataset, index_name: &str) -> Re Ok(()) } +async fn prewarm_colliding_segments( + indices: Vec, + named_field_id: i32, + open_and_prewarm: F, +) -> Result> +where + F: Fn(i32, Uuid) -> Fut, + Fut: Future>, +{ + // Both opening and prewarming can perform remote I/O. Keep the entire + // segment operation concurrent, with the SDK's fail-fast cancellation. + let segments = futures::future::try_join_all(indices.into_iter().map(|metadata| { + let open_and_prewarm = &open_and_prewarm; + async move { + let Some(&field_id) = metadata.fields.first() else { + return Ok::<_, lance_core::Error>(Some(metadata.uuid)); + }; + if field_id == named_field_id || !open_and_prewarm(field_id, metadata.uuid).await? { + Ok(Some(metadata.uuid)) + } else { + Ok(None) + } + } + })) + .await?; + Ok(segments.into_iter().flatten().collect()) +} + +#[cfg(test)] +#[path = "index_tests.rs"] +mod tests; + /// Drop an index by name. #[unsafe(no_mangle)] pub unsafe extern "C" fn lance_dataset_drop_index( diff --git a/src/index_tests.rs b/src/index_tests.rs new file mode 100644 index 0000000..654e3e8 --- /dev/null +++ b/src/index_tests.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::cell::Cell; + +use futures::{FutureExt, channel::oneshot, poll}; + +use super::*; + +fn segment(id: u128, field_id: Option) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::from_u128(id), + fields: field_id.into_iter().collect(), + covering_fields: Vec::new(), + name: "id".into(), + dataset_version: 1, + fragment_bitmap: None, + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } +} + +#[test] +fn prewarm_colliding_segments_opens_and_warms_concurrently() { + futures::executor::block_on(async { + const SEGMENTS: usize = 8; + let opening = Cell::new(0); + let warming = Cell::new(0); + let (release_open, open_gate) = oneshot::channel::<()>(); + let (release_warm, warm_gate) = oneshot::channel::<()>(); + let open_gate = open_gate.shared(); + let warm_gate = warm_gate.shared(); + let indices = (0..SEGMENTS) + .map(|id| segment(id as u128, Some(1))) + .collect(); + let mut prewarm = Box::pin(prewarm_colliding_segments(indices, 0, |field_id, _| { + assert_eq!(field_id, 1); + let open_gate = open_gate.clone(); + let warm_gate = warm_gate.clone(); + let opening = &opening; + let warming = &warming; + async move { + opening.set(opening.get() + 1); + open_gate.await.unwrap(); + warming.set(warming.get() + 1); + warm_gate.await.unwrap(); + Ok(true) + } + })); + + // Gates make this independent of storage speed and scheduler timing: + // every open must start before any open is allowed to finish. + assert!(poll!(&mut prewarm).is_pending()); + assert_eq!(opening.get(), SEGMENTS); + assert_eq!(warming.get(), 0); + release_open.send(()).unwrap(); + assert!(poll!(&mut prewarm).is_pending()); + assert_eq!(warming.get(), SEGMENTS); + release_warm.send(()).unwrap(); + assert!(prewarm.await.unwrap().is_empty()); + }); +} + +struct Dropped<'a>(&'a Cell); + +impl Drop for Dropped<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } +} + +#[test] +fn prewarm_colliding_segments_cancels_pending_work_on_error() { + for fail_during_open in [true, false] { + futures::executor::block_on(async { + const SEGMENTS: usize = 8; + let started = Cell::new(0); + let dropped = Cell::new(0); + let (release, gate) = oneshot::channel::<()>(); + let gate = gate.shared(); + let indices = (0..SEGMENTS) + .map(|id| segment(id as u128, Some(1))) + .collect(); + let mut prewarm = Box::pin(prewarm_colliding_segments(indices, 0, |_, uuid| { + let gate = gate.clone(); + let started = &started; + let dropped = &dropped; + async move { + let _guard = Dropped(dropped); + started.set(started.get() + 1); + gate.await.unwrap(); + if !fail_during_open { + // Allow all opens to finish before testing a prewarm failure. + futures::pending!(); + } + if uuid == Uuid::from_u128(3) { + return Err(lance_core::Error::invalid_input_source( + "segment I/O failed".into(), + )); + } + futures::future::pending::>().await + } + })); + + assert!(poll!(&mut prewarm).is_pending()); + assert_eq!(started.get(), SEGMENTS); + release.send(()).unwrap(); + if !fail_during_open { + assert!(poll!(&mut prewarm).is_pending()); + } + let result = poll!(&mut prewarm); + assert!( + result.is_ready(), + "a failed segment must not wait for peers" + ); + let err = match result { + std::task::Poll::Ready(Err(err)) => err, + other => panic!("expected the segment error, got {other:?}"), + }; + assert!(err.to_string().contains("segment I/O failed")); + assert_eq!(dropped.get(), SEGMENTS, "pending work must be cancelled"); + }); + } +} + +#[test] +fn prewarm_colliding_segments_preserves_sdk_fallback() { + futures::executor::block_on(async { + let opened = Cell::new(0); + let indices = vec![ + segment(0, None), + segment(1, Some(0)), + segment(2, Some(1)), + segment(3, Some(1)), + ]; + let sdk_segments = prewarm_colliding_segments(indices, 0, |field_id, uuid| { + assert_eq!(field_id, 1); + opened.set(opened.get() + 1); + // Segment 2 is scalar; segment 3 is vector. + futures::future::ready(Ok(uuid == Uuid::from_u128(3))) + }) + .await + .unwrap(); + assert_eq!(opened.get(), 2); + assert_eq!( + sdk_segments, + [0, 1, 2].map(Uuid::from_u128), + "scalar and unresolved segments must retain SDK validation" + ); + }); +}