fix(python): scope Instructor compatibility shim

python/python/lancedb/embeddings/instructor.py

@@ -3,6 +3,7 @@
 
 
 from typing import List
+from urllib.parse import unquote, urlparse
 
 import numpy as np
 
@@ -125,9 +126,20 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
 
     @weak_lru(maxsize=1)
     def get_model(self):
-        instructor_embedding = attempt_import_or_raise(
-            "InstructorEmbedding", "InstructorEmbedding"
-        )
+        huggingface_hub = attempt_import_or_raise("huggingface_hub", "huggingface-hub")
+        missing = object()
+        original_cached_download = getattr(huggingface_hub, "cached_download", missing)
+        if original_cached_download is missing:
+            huggingface_hub.cached_download = _cached_download(huggingface_hub)
+
+        try:
+            instructor_embedding = attempt_import_or_raise(
+                "InstructorEmbedding", "InstructorEmbedding"
+            )
+        finally:
+            if original_cached_download is missing:
+                del huggingface_hub.cached_download
+
         torch = attempt_import_or_raise("torch", "torch")
 
         model = instructor_embedding.INSTRUCTOR(self.name)
@@ -140,3 +152,44 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
                 model, {torch.nn.Linear}, dtype=torch.qint8
             )
         return model
+
+
+def _cached_download(huggingface_hub):
+    """Provide the legacy download API used by sentence-transformers 2.2.x."""
+
+    def cached_download(
+        *,
+        url,
+        cache_dir=None,
+        force_filename=None,
+        library_name=None,
+        library_version=None,
+        user_agent=None,
+        use_auth_token=None,
+        **_,
+    ):
+        path = urlparse(url).path.lstrip("/")
+        try:
+            repo_id, resolved_path = path.split("/resolve/", maxsplit=1)
+            revision, filename = resolved_path.split("/", maxsplit=1)
+        except ValueError as err:
+            raise ValueError(f"Unsupported Hugging Face Hub URL: {url}") from err
+
+        repo_id = unquote(repo_id)
+        revision = unquote(revision)
+        filename = unquote(filename)
+        # sentence-transformers derives force_filename from this Hub path with
+        # os.path.join. Using the URL path beneath local_dir produces the same
+        # local destination without sending Windows separators to the Hub.
+        return huggingface_hub.hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+            local_dir=cache_dir,
+            library_name=library_name,
+            library_version=library_version,
+            user_agent=user_agent,
+            token=use_auth_token,
+        )
+
+    return cached_download

python/python/tests/test_embeddings.py

@@ -1,8 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+import ntpath
 import os
 import pickle
+import sys
+from types import ModuleType
 from typing import List, Optional, Union
 from unittest.mock import MagicMock, patch
 
@@ -522,6 +525,59 @@ def test_embedding_function_safe_model_dump(embedding_type):
             )
 
 
+def test_instructor_embedding_supports_huggingface_hub_without_cached_download(
+    tmp_path, monkeypatch
+):
+    from lancedb.embeddings.instructor import InstructorEmbeddingFunction
+
+    hub_download = MagicMock(return_value="/cache/1_Pooling/config.json")
+    huggingface_hub = ModuleType("huggingface_hub")
+    huggingface_hub.hf_hub_download = hub_download
+    torch = ModuleType("torch")
+    monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.delitem(sys.modules, "InstructorEmbedding", raising=False)
+    monkeypatch.syspath_prepend(str(tmp_path))
+
+    (tmp_path / "InstructorEmbedding.py").write_text(
+        "from huggingface_hub import cached_download\n\n"
+        "class INSTRUCTOR:\n"
+        "    def __init__(self, name):\n"
+        "        self.name = name\n"
+    )
+
+    embedding = InstructorEmbeddingFunction.create(show_progress_bar=False)
+    instructor_model = embedding.get_model()
+
+    assert instructor_model.name == "hkunlp/instructor-base"
+    assert not hasattr(huggingface_hub, "cached_download")
+
+    instructor_embedding = sys.modules["InstructorEmbedding"]
+    path = instructor_embedding.cached_download(
+        url=(
+            "https://huggingface.co/hkunlp/instructor-base/resolve/abc123/"
+            "1_Pooling/config.json"
+        ),
+        cache_dir="/cache",
+        force_filename=ntpath.join("1_Pooling", "config.json"),
+        library_name="sentence-transformers",
+        library_version="2.2.2",
+        use_auth_token="token",
+    )
+
+    assert path == "/cache/1_Pooling/config.json"
+    hub_download.assert_called_once_with(
+        repo_id="hkunlp/instructor-base",
+        filename="1_Pooling/config.json",
+        revision="abc123",
+        local_dir="/cache",
+        library_name="sentence-transformers",
+        library_version="2.2.2",
+        user_agent=None,
+        token="token",
+    )
+
+
 @patch("time.sleep")
 def test_retry(mock_sleep):
     test_function = MagicMock(side_effect=[Exception] * 9 + ["result"])

rust/lancedb/src/table/create_index.rs

@@ -416,7 +416,7 @@ mod tests {
     use arrow_array::record_batch;
     use arrow_array::{
         Array, ArrayRef, BinaryArray, BooleanArray, FixedSizeListArray, Float32Array, Int32Array,
-        LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StructArray, UInt64Array,
+        LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StructArray,
     };
     use arrow_data::ArrayDataBuilder;
     use arrow_schema::{DataType, Field, Schema};
@@ -436,55 +436,6 @@ mod tests {
     use crate::table::optimize::{CompactionOptions, OptimizeAction};
     use lance_index::scalar::FullTextSearchQuery;
 
-    struct OrderPreservingShuffler;
-
-    struct SinglePartitionReader {
-        batches: Vec<RecordBatch>,
-    }
-
-    #[async_trait::async_trait]
-    impl lance_index::vector::v3::shuffler::Shuffler for OrderPreservingShuffler {
-        async fn shuffle(
-            &self,
-            data: Box<dyn lance_io::stream::RecordBatchStream + Unpin + 'static>,
-        ) -> lance_core::Result<Box<dyn lance_index::vector::v3::shuffler::ShuffleReader>> {
-            Ok(Box::new(SinglePartitionReader {
-                batches: data.try_collect().await?,
-            }))
-        }
-    }
-
-    #[async_trait::async_trait]
-    impl lance_index::vector::v3::shuffler::ShuffleReader for SinglePartitionReader {
-        async fn read_partition(
-            &self,
-            partition_id: usize,
-        ) -> lance_core::Result<
-            Option<Box<dyn lance_io::stream::RecordBatchStream + Unpin + 'static>>,
-        > {
-            if partition_id != 0 || self.batches.is_empty() {
-                return Ok(None);
-            }
-            let schema = self.batches[0].schema();
-            let stream = futures::stream::iter(self.batches.clone().into_iter().map(Ok));
-            Ok(Some(Box::new(
-                lance_io::stream::RecordBatchStreamAdapter::new(schema, stream),
-            )))
-        }
-
-        fn partition_size(&self, partition_id: usize) -> lance_core::Result<usize> {
-            Ok(if partition_id == 0 {
-                self.batches.iter().map(RecordBatch::num_rows).sum()
-            } else {
-                0
-            })
-        }
-
-        fn total_loss(&self) -> Option<f64> {
-            Some(0.0)
-        }
-    }
-
     fn create_fixed_size_list<T: Array>(
         values: T,
         list_size: i32,
@@ -889,135 +840,6 @@ mod tests {
         assert_eq!(stats.distance_type, Some(crate::DistanceType::L2));
     }
 
-    #[tokio::test]
-    async fn test_ivf_hnsw_sq_merge_preserves_shuffled_row_order() {
-        use lance::index::vector::builder::IvfIndexBuilder;
-        use lance_core::{ROW_ID, ROW_ID_FIELD, cache::LanceCache};
-        use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
-        use lance_file::reader::{FileReader, FileReaderOptions};
-        use lance_index::INDEX_AUXILIARY_FILE_NAME;
-        use lance_index::vector::hnsw::{HNSW, builder::HnswBuildParams};
-        use lance_index::vector::ivf::IvfBuildParams;
-        use lance_index::vector::sq::{ScalarQuantizer, builder::SQBuildParams};
-        use lance_io::ReadBatchParams;
-        use lance_io::object_store::ObjectStore;
-        use lance_io::scheduler::{ScanScheduler, SchedulerConfig};
-        use lance_io::stream::RecordBatchStreamAdapter;
-        use lance_io::utils::CachedFileSize;
-
-        const NUM_ROWS: usize = 64;
-        const DIMENSION: i32 = 16;
-
-        let tmp_dir = tempdir().unwrap();
-        let conn = connect(tmp_dir.path().to_str().unwrap())
-            .execute()
-            .await
-            .unwrap();
-        let vector_field = Field::new(
-            "embeddings",
-            DataType::FixedSizeList(
-                Arc::new(Field::new("item", DataType::Float32, true)),
-                DIMENSION,
-            ),
-            false,
-        );
-        let values = Float32Array::from_iter_values(
-            (0..NUM_ROWS * DIMENSION as usize).map(|value| (value % 97) as f32),
-        );
-        let vectors = Arc::new(create_fixed_size_list(values, DIMENSION).unwrap());
-        let table_batch = RecordBatch::try_new(
-            Arc::new(Schema::new(vec![vector_field.clone()])),
-            vec![vectors.clone()],
-        )
-        .unwrap();
-        let table = conn
-            .create_table("test", table_batch)
-            .execute()
-            .await
-            .unwrap();
-
-        let dataset_guard = table.as_native().unwrap().dataset.get().await.unwrap();
-        let dataset = (*dataset_guard).clone();
-        drop(dataset_guard);
-
-        // Descending row IDs are a small sentinel for the removed merge-time
-        // sort. The old sort reordered this batch and used FixedSizeList::take,
-        // which overflowed for the large partitions reported in #3126.
-        let expected_row_ids = (0..NUM_ROWS as u64).rev().collect::<Vec<_>>();
-        let input_batch = RecordBatch::try_new(
-            Arc::new(Schema::new(vec![ROW_ID_FIELD.clone(), vector_field])),
-            vec![
-                Arc::new(UInt64Array::from(expected_row_ids.clone())),
-                vectors,
-            ],
-        )
-        .unwrap();
-        let input = RecordBatchStreamAdapter::new(
-            input_batch.schema(),
-            futures::stream::iter(vec![Ok(input_batch)]),
-        );
-
-        let index_dir = dataset.indices_dir().join(uuid::Uuid::new_v4().to_string());
-        let mut builder = IvfIndexBuilder::<HNSW, ScalarQuantizer>::new(
-            dataset,
-            "embeddings".to_string(),
-            index_dir.clone(),
-            lance_linalg::distance::DistanceType::L2,
-            Box::new(OrderPreservingShuffler),
-            Some(IvfBuildParams::new(1)),
-            Some(SQBuildParams::default()),
-            HnswBuildParams::default(),
-            None,
-        )
-        .unwrap();
-        builder.shuffle_data_input(Some(input));
-        builder.build().await.unwrap();
-
-        let object_store = Arc::new(ObjectStore::local());
-        let scheduler = ScanScheduler::new(object_store, SchedulerConfig::default_for_testing());
-        let auxiliary_path = index_dir.join(INDEX_AUXILIARY_FILE_NAME);
-        let reader = FileReader::try_open(
-            scheduler
-                .open_file(&auxiliary_path, &CachedFileSize::unknown())
-                .await
-                .unwrap(),
-            None,
-            Arc::<DecoderPlugins>::default(),
-            &LanceCache::no_cache(),
-            FileReaderOptions::default(),
-        )
-        .await
-        .unwrap();
-        let batches = reader
-            .read_stream(
-                ReadBatchParams::RangeFull,
-                u32::MAX,
-                1,
-                FilterExpression::no_filter(),
-            )
-            .await
-            .unwrap()
-            .try_collect::<Vec<_>>()
-            .await
-            .unwrap();
-        let stored_row_ids = batches
-            .iter()
-            .flat_map(|batch| {
-                batch
-                    .column_by_name(ROW_ID)
-                    .unwrap()
-                    .as_any()
-                    .downcast_ref::<UInt64Array>()
-                    .unwrap()
-                    .values()
-                    .iter()
-                    .copied()
-            })
-            .collect::<Vec<_>>();
-
-        assert_eq!(stored_row_ids, expected_row_ids);
-    }
-
     #[tokio::test]
     async fn test_create_index_ivf_hnsw_pq() {
         use std::iter::repeat_with;