feat: Qdrant Search Index

README.md

@@ -45,7 +45,7 @@
 | 🔄 **Smart caching** | Never wait for your data to process twice — cached results are automatically reused |
 | 📊 **Multi-framework interoperability** | Native conversion to/from NumPy, Pandas, Polars, Arrow, PyTorch, TensorFlow, JAX, and Spark |
 | 🏎️ **Multi-processing** | Fast parallel data processing with `map(num_proc=N)` |
-| 🔎 **Search & index** | Built-in FAISS and Elasticsearch index support for similarity search |
+| 🔎 **Search & index** | Built-in FAISS, Elasticsearch, and Qdrant index support for similarity search |
 | 📦 **JSON type** | Flexible JSON/structured data support with `Json()` feature type |
 
 # Installation

docs/source/_redirects.yml

@@ -6,7 +6,8 @@ dataset_streaming: stream
 torch_tensorflow: use_dataset
 splits: loading#slice-splits
 processing: process
-faiss_and_ea: faiss_es
+faiss_es: faiss_es_qdrant
+faiss_and_ea: faiss_es_qdrant
 features: about_dataset_features
 exploring: access
 package_reference/logging_methods: package_reference/utilities

docs/source/_toctree.yml

@@ -50,7 +50,7 @@
       title: Cache management
     - local: filesystems
       title: Cloud storage
-    - local: faiss_es
+    - local: faiss_es_qdrant
       title: Search index
     - local: cli
       title: CLI

docs/source/faiss_es.mdx → docs/source/faiss_es_qdrant.mdx

@@ -1,6 +1,6 @@
 # Search index
 
-[FAISS](https://github.com/facebookresearch/faiss) and [Elasticsearch](https://www.elastic.co/elasticsearch/) enables searching for examples in a dataset. This can be useful when you want to retrieve specific examples from a dataset that are relevant to your NLP task. For example, if you are working on an Open Domain Question Answering task, you may want to only return examples that are relevant to answering your question.
+[FAISS](https://github.com/facebookresearch/faiss), [Elasticsearch](https://www.elastic.co/elasticsearch/), and [Qdrant](https://qdrant.tech/) enable searching for examples in a dataset. This can be useful when you want to retrieve specific examples from a dataset that are relevant to your NLP task. For example, if you are working on an Open Domain Question Answering task, you may want to only return examples that are relevant to answering your question.
 
 This guide will show you how to build an index for your dataset that will allow you to search it.
 
@@ -131,3 +131,45 @@ For more advanced Elasticsearch usage, you can specify your own configuration wi
 >>> es_index_name = "hf_squad_context"  # name of the index in Elasticsearch
 >>> squad.add_elasticsearch_index("context", es_client=es_client, es_config=es_config, es_index_name=es_index_name)
 ```
+
+## Qdrant
+
+[Qdrant](https://qdrant.tech/) provides persistent vector search with optional payload filtering. Install the optional dependency first:
+
+```bash
+pip install "datasets[qdrant]"
+```
+
+Connect to a Qdrant server using its URL. The following example stores the `language` column as payload so that it can be used as a search filter:
+
+```py
+>>> from qdrant_client import QdrantClient, models
+>>> client = QdrantClient(url="http://localhost:6333")
+>>> ds_with_embeddings = ds_with_embeddings.add_column(
+...     "language", ["en"] * len(ds_with_embeddings)
+... )
+>>> ds_with_embeddings.add_qdrant_index(
+...     column="embeddings",
+...     payload_indexes={"language": "keyword"},
+...     collection_name="documents",
+...     qdrant_client=client,
+... )
+>>> query_filter = models.Filter(
+...     must=[models.FieldCondition(key="language", match=models.MatchValue(value="en"))]
+... )
+>>> scores, examples = ds_with_embeddings.get_nearest_examples(
+...     "embeddings", question_embedding, k=10, query_filter=query_filter
+... )
+```
+
+Fields in `payload_indexes` are uploaded as payload and indexed before the vectors are added. Use `payload_columns` for additional fields that should be stored without an index. Only index fields used in filters, because payload indexes consume additional memory.
+
+Dataset row numbers are used as Qdrant point IDs, so results map back to examples without copying the full dataset into Qdrant. [`Dataset.load_qdrant_index`] only attaches collections created by [`Dataset.add_qdrant_index`], because it relies on this row-ID mapping. To attach one later:
+
+```py
+>>> ds.load_qdrant_index(
+...     "embeddings",
+...     qdrant_client=QdrantClient(url="http://localhost:6333"),
+...     collection_name="documents",
+... )
+```

docs/source/package_reference/main_classes.mdx

@@ -72,6 +72,8 @@ The base class [`Dataset`] implements a Dataset backed by an Apache Arrow table.
     - load_faiss_index
     - add_elasticsearch_index
     - load_elasticsearch_index
+    - add_qdrant_index
+    - load_qdrant_index
     - list_indexes
     - get_index
     - drop_index

setup.py

@@ -145,6 +145,8 @@
     "Pillow>=9.4.0",  # When PIL.Image.ExifTags was introduced
 ]
 
+QDRANT_REQUIRE = ["qdrant-client>=1.10.0,<2.0.0"]
+
 MESH_REQUIRE = [
     "trimesh>=4.10.0",
 ]
@@ -170,6 +172,7 @@
     "aiohttp",
     "elasticsearch>=7.17.12,<8.0.0",  # 8.0 asks users to provide hosts or cloud_id when instantiating ElasticSearch(); 7.9.1 has legacy numpy.float_ which was fixed in https://github.com/elastic/elasticsearch-py/pull/2551.
     "faiss-cpu>=1.8.0.post1",  # Pins numpy < 2
+    *QDRANT_REQUIRE,
     "h5py",
     "pylance",
     "pyiceberg[sql-sqlite,pyarrow]",
@@ -230,6 +233,7 @@
     "tensorflow_gpu": ["tensorflow>=2.6.0"],
     "torch": ["torch"],
     "jax": ["jax>=0.3.14", "jaxlib>=0.3.14"],
+    "qdrant": QDRANT_REQUIRE,
     "streaming": [],  # for backward compatibility
     "dev": TESTS_REQUIRE + QUALITY_REQUIRE + DOCS_REQUIRE,
     "tests": TESTS_REQUIRE,

src/datasets/search.py

@@ -25,9 +25,16 @@
 
     except ImportError:
         pass
+    try:
+        from qdrant_client import QdrantClient  # noqa: F401
+        from qdrant_client.models import Distance, PayloadSchemaType  # noqa: F401
+
+    except ImportError:
+        pass
 
 _has_elasticsearch = importlib.util.find_spec("elasticsearch") is not None
 _has_faiss = importlib.util.find_spec("faiss") is not None
+_has_qdrant = importlib.util.find_spec("qdrant_client") is not None
 
 
 logger = logging.get_logger(__name__)
@@ -414,6 +421,143 @@ def load(
         return faiss_index
 
 
+class QdrantIndex(BaseIndex):
+    """Dense vector index backed by an existing Qdrant client."""
+
+    @staticmethod
+    def _as_qdrant_payload(value):
+        if isinstance(value, np.generic):
+            return QdrantIndex._as_qdrant_payload(value.item())
+        if isinstance(value, np.ndarray):
+            return QdrantIndex._as_qdrant_payload(value.tolist())
+        if isinstance(value, dict):
+            return {key: QdrantIndex._as_qdrant_payload(item) for key, item in value.items()}
+        if isinstance(value, (list, tuple)):
+            return [QdrantIndex._as_qdrant_payload(item) for item in value]
+        return value
+
+    def __init__(
+        self,
+        qdrant_client: "QdrantClient",
+        collection_name: str,
+        distance: Union[str, "Distance"] = "Cosine",
+    ):
+        if not _has_qdrant:
+            raise ImportError("Install Qdrant support with `pip install datasets[qdrant]`.")
+        self.qdrant_client = qdrant_client
+        self.collection_name = collection_name
+        self.distance = distance
+
+    def add_vectors(
+        self,
+        dataset: "Dataset",
+        column: str,
+        payload_columns: Optional[list[str]] = None,
+        payload_indexes: Optional[dict[str, Union[str, "PayloadSchemaType"]]] = None,
+        batch_size: int = 64,
+        **upload_kwargs,
+    ):
+        """Create a collection and upload dataset rows as Qdrant points."""
+        from qdrant_client import models
+
+        if not isinstance(dataset.features[column], List):
+            raise ValueError(
+                f"Wrong feature type for column '{column}'. Expected 1d array, got {dataset.features[column]}"
+            )
+        if not len(dataset):
+            raise ValueError("Cannot create a Qdrant index from an empty dataset.")
+        if batch_size <= 0:
+            raise ValueError("batch_size must be greater than 0.")
+        payload_indexes = payload_indexes or {}
+        payload_columns = set(payload_columns or []).union(payload_indexes)
+        missing_columns = payload_columns - set(dataset.column_names)
+        if missing_columns:
+            raise ValueError(f"Payload columns not found in the dataset: {sorted(missing_columns)}")
+        payload_indexes = {
+            field_name: models.PayloadSchemaType(field_schema) if isinstance(field_schema, str) else field_schema
+            for field_name, field_schema in payload_indexes.items()
+        }
+        self.qdrant_client.create_collection(
+            collection_name=self.collection_name,
+            vectors_config=models.VectorParams(
+                size=len(dataset[0][column]),
+                distance=(
+                    models.Distance(self.distance.capitalize()) if isinstance(self.distance, str) else self.distance
+                ),
+            ),
+        )
+        for field_name, field_schema in payload_indexes.items():
+            self.qdrant_client.create_payload_index(
+                collection_name=self.collection_name,
+                field_name=field_name,
+                field_schema=field_schema,
+                wait=True,
+            )
+
+        def point_generator():
+            for offset in range(0, len(dataset), batch_size):
+                batch = dataset[offset : offset + batch_size]
+                for i, vector in enumerate(batch[column]):
+                    yield models.PointStruct(
+                        id=offset + i,
+                        vector=vector,
+                        payload={name: self._as_qdrant_payload(batch[name][i]) for name in payload_columns} or None,
+                    )
+
+        logger.info(f"Adding {len(dataset)} vectors to the Qdrant collection {self.collection_name}")
+        upload_kwargs.setdefault("wait", True)
+        self.qdrant_client.upload_points(
+            collection_name=self.collection_name,
+            points=hf_tqdm(point_generator(), unit="points", total=len(dataset)),
+            batch_size=batch_size,
+            **upload_kwargs,
+        )
+
+    def search(self, query: np.array, k=10, **kwargs) -> SearchResults:
+        """Find the nearest dataset rows to a query vector."""
+        query = np.asarray(query, dtype=np.float32)
+        if query.ndim != 1 and (query.ndim != 2 or query.shape[0] != 1):
+            raise ValueError("Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)")
+
+        kwargs.setdefault("with_payload", False)
+        kwargs.setdefault("with_vectors", False)
+        response = self.qdrant_client.query_points(
+            collection_name=self.collection_name,
+            query=query.reshape(-1).tolist(),
+            limit=k,
+            **kwargs,
+        )
+        return SearchResults([point.score for point in response.points], [int(point.id) for point in response.points])
+
+    def search_batch(self, queries: np.array, k=10, **kwargs) -> BatchedSearchResults:
+        from qdrant_client import models
+
+        queries = np.asarray(queries, dtype=np.float32)
+        if queries.ndim != 2:
+            raise ValueError("Shape of query must be 2D")
+
+        batch_kwargs = {name: kwargs.pop(name) for name in ("consistency", "timeout") if name in kwargs}
+        for query_name, request_name in {
+            "query_filter": "filter",
+            "search_params": "params",
+            "shard_key_selector": "shard_key",
+            "with_vectors": "with_vector",
+        }.items():
+            if query_name in kwargs:
+                kwargs[request_name] = kwargs.pop(query_name)
+        kwargs.setdefault("with_payload", False)
+        kwargs.setdefault("with_vector", False)
+        responses = self.qdrant_client.query_batch_points(
+            collection_name=self.collection_name,
+            requests=[models.QueryRequest(query=query.tolist(), limit=k, **kwargs) for query in queries],
+            **batch_kwargs,
+        )
+        return BatchedSearchResults(
+            total_scores=[[point.score for point in response.points] for response in responses],
+            total_indices=[[int(point.id) for point in response.points] for response in responses],
+        )
+
+
 class IndexableMixin:
     """Add indexing features to `datasets.Dataset`"""
 
@@ -432,7 +576,8 @@ def is_index_initialized(self, index_name: str) -> bool:
     def _check_index_is_initialized(self, index_name: str):
         if not self.is_index_initialized(index_name):
             raise MissingIndex(
-                f"Index with index_name '{index_name}' not initialized yet. Please make sure that you call `add_faiss_index` or `add_elasticsearch_index` first."
+                f"Index with index_name '{index_name}' not initialized yet. Please make sure that you call "
+                "`add_faiss_index`, `add_elasticsearch_index`, or `add_qdrant_index` first."
             )
 
     def list_indexes(self) -> list[str]:
@@ -681,6 +826,76 @@ def load_elasticsearch_index(
             host=host, port=port, es_client=es_client, es_index_name=es_index_name, es_index_config=es_index_config
         )
 
+    def add_qdrant_index(
+        self,
+        column: str,
+        qdrant_client: "QdrantClient",
+        collection_name: str,
+        index_name: Optional[str] = None,
+        distance: Union[str, "Distance"] = "Cosine",
+        payload_columns: Optional[list[str]] = None,
+        payload_indexes: Optional[dict[str, Union[str, "PayloadSchemaType"]]] = None,
+        batch_size: int = 64,
+        **upload_kwargs,
+    ):
+        """Add a dense vector index backed by Qdrant.
+
+        Args:
+            column (`str`): Vector column to index.
+            qdrant_client (`qdrant_client.QdrantClient`): Client configured with a Qdrant URL.
+            collection_name (`str`): Qdrant collection to create.
+            index_name (`str`, *optional*): Dataset index name. Defaults to `column`.
+            distance (`str` or `qdrant_client.models.Distance`): Vector distance.
+            payload_columns (`List[str]`, *optional*): Columns stored as filterable Qdrant payload.
+            payload_indexes (`Dict[str, str or qdrant_client.models.PayloadSchemaType]`, *optional*):
+                Mapping of payload field names to Qdrant index schemas..
+            batch_size (`int`): Points per upload batch.
+            **upload_kwargs: Additional arguments for `QdrantClient.upload_points`.
+        """
+        qdrant_index = QdrantIndex(
+            qdrant_client=qdrant_client,
+            collection_name=collection_name,
+            distance=distance,
+        )
+        qdrant_index.add_vectors(
+            self,
+            column=column,
+            payload_columns=payload_columns,
+            payload_indexes=payload_indexes,
+            batch_size=batch_size,
+            **upload_kwargs,
+        )
+        self._indexes[index_name if index_name is not None else column] = qdrant_index
+        return self
+
+    def load_qdrant_index(
+        self,
+        index_name: str,
+        qdrant_client: "QdrantClient",
+        collection_name: str,
+    ):
+        """Attach a Qdrant collection previously created by `add_qdrant_index`.
+
+        The collection must contain one point per dataset row, with point IDs equal to
+        the zero-based dataset row numbers.
+
+        Args:
+            index_name (`str`): Dataset index name.
+            qdrant_client (`qdrant_client.QdrantClient`): Client connected to the collection.
+            collection_name (`str`): Existing Qdrant collection name.
+        """
+        index = QdrantIndex(
+            qdrant_client=qdrant_client,
+            collection_name=collection_name,
+        )
+        points_count = index.qdrant_client.count(collection_name=collection_name, exact=True).count
+        if points_count != len(self):
+            raise ValueError(
+                f"Index size should match Dataset size, but Qdrant collection '{collection_name}' has "
+                f"{points_count} points while the dataset has {len(self)} examples."
+            )
+        self._indexes[index_name] = index
+
     def drop_index(self, index_name: str):
         """Drop the index with the specified column.