fix: resolve vector prewarm name collisions from segment metadata

src/index.rs

@@ -7,13 +7,17 @@
 //! 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};
-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 lance_table::format::IndexMetadata;
+use uuid::Uuid;
 
 use crate::dataset::LanceDataset;
 use crate::error::ffi_try;
@@ -297,10 +301,76 @@ 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 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
+            .prewarm_index_segments(index_name, &sdk_segments)
+            .await?;
+    }
+    Ok(())
+}
+
+async fn prewarm_colliding_segments<F, Fut>(
+    indices: Vec<IndexMetadata>,
+    named_field_id: i32,
+    open_and_prewarm: F,
+) -> Result<Vec<Uuid>>
+where
+    F: Fn(i32, Uuid) -> Fut,
+    Fut: Future<Output = Result<bool>>,
+{
+    // 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(

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<i32>) -> 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<usize>);
+
+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::<Result<bool>>().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"
+        );
+    });
+}

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, &params)
-                    .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);