From 037bf5c54dc1f21eed426b97cee9f18d6b35681c Mon Sep 17 00:00:00 2001 From: Anush008 Date: Fri, 31 Jul 2026 18:55:17 +0530 Subject: [PATCH] feat: Qdrant search index --- README.md | 2 +- docs/source/_redirects.yml | 3 +- docs/source/_toctree.yml | 2 +- .../{faiss_es.mdx => faiss_es_qdrant.mdx} | 44 +++- .../source/package_reference/main_classes.mdx | 2 + setup.py | 4 + src/datasets/search.py | 217 +++++++++++++++++- tests/test_search.py | 64 +++++- tests/utils.py | 1 + 9 files changed, 332 insertions(+), 7 deletions(-) rename docs/source/{faiss_es.mdx => faiss_es_qdrant.mdx} (70%) diff --git a/README.md b/README.md index d96cbda38ab..1f9effb7430 100644 --- a/README.md +++ b/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 diff --git a/docs/source/_redirects.yml b/docs/source/_redirects.yml index 606a00ea1cc..c6a093e73dd 100644 --- a/docs/source/_redirects.yml +++ b/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 diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 276fe36c444..babfd1fa8c3 100644 --- a/docs/source/_toctree.yml +++ b/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 diff --git a/docs/source/faiss_es.mdx b/docs/source/faiss_es_qdrant.mdx similarity index 70% rename from docs/source/faiss_es.mdx rename to docs/source/faiss_es_qdrant.mdx index 635051744de..28ae98ae63c 100644 --- a/docs/source/faiss_es.mdx +++ b/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", +... ) +``` diff --git a/docs/source/package_reference/main_classes.mdx b/docs/source/package_reference/main_classes.mdx index 5d3ce6588b5..c889e898700 100644 --- a/docs/source/package_reference/main_classes.mdx +++ b/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 diff --git a/setup.py b/setup.py index c536574a15e..c51110cf73b 100644 --- a/setup.py +++ b/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, diff --git a/src/datasets/search.py b/src/datasets/search.py index cb994d24e41..eae19413f6e 100644 --- a/src/datasets/search.py +++ b/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. diff --git a/tests/test_search.py b/tests/test_search.py index 142c92522ea..9b43e2051d2 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -2,7 +2,7 @@ import tempfile from functools import partial from unittest import TestCase -from unittest.mock import patch +from unittest.mock import Mock, patch import numpy as np import pytest @@ -10,7 +10,7 @@ from datasets.arrow_dataset import Dataset from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex -from .utils import require_elasticsearch, require_faiss +from .utils import require_elasticsearch, require_faiss, require_qdrant pytestmark = pytest.mark.integration @@ -242,3 +242,63 @@ def test_elasticsearch(self): best_indices = [indices[0] for indices in total_indices] self.assertGreater(np.min(best_scores), 0) self.assertListEqual([1, 1, 1], best_indices) + + +@require_qdrant +class QdrantIndexTest(TestCase): + def test_dataset_index(self): + from qdrant_client import models + + client = Mock() + uploaded_points = [] + client.upload_points.side_effect = lambda **kwargs: uploaded_points.extend(kwargs["points"]) + client.query_points.return_value = Mock(points=[Mock(score=1.0, id=1)]) + client.query_batch_points.return_value = [ + Mock(points=[Mock(score=1.0, id=1)]), + Mock(points=[Mock(score=1.0, id=0)]), + ] + client.count.return_value = Mock(count=2) + dset = Dataset.from_dict( + { + "text": ["zero", "one"], + "category": ["even", "odd"], + "embeddings": np.eye(2, dtype=np.float32), + } + ) + + self.assertIs( + dset.add_qdrant_index( + "embeddings", + qdrant_client=client, + collection_name="documents", + distance="Dot", + payload_indexes={"category": "keyword"}, + batch_size=1, + ), + dset, + ) + self.assertEqual( + client.create_payload_index.call_args.kwargs, + { + "collection_name": "documents", + "field_name": "category", + "field_schema": models.PayloadSchemaType.KEYWORD, + "wait": True, + }, + ) + method_names = [call[0] for call in client.method_calls] + self.assertLess(method_names.index("create_payload_index"), method_names.index("upload_points")) + self.assertEqual([point.id for point in uploaded_points], [0, 1]) + self.assertEqual(dset.get_nearest_examples("embeddings", [0, 1], k=1).examples["text"], ["one"]) + query_filter = models.Filter(must=[]) + self.assertEqual( + dset.search_batch( + "embeddings", np.eye(2, dtype=np.float32)[::-1], k=1, query_filter=query_filter + ).total_indices, + [[1], [0]], + ) + Dataset.from_dict(dset.to_dict()).load_qdrant_index( + "embeddings", qdrant_client=client, collection_name="documents" + ) + self.assertEqual([point.payload for point in uploaded_points], [{"category": "even"}, {"category": "odd"}]) + self.assertEqual(client.query_batch_points.call_args.kwargs["requests"][0].filter, query_filter) diff --git a/tests/utils.py b/tests/utils.py index ed20acca96c..789386381de 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -62,6 +62,7 @@ def parse_flag_from_env(key, default=False): require_faiss = pytest.mark.skipif(find_spec("faiss") is None or sys.platform == "win32", reason="test requires faiss") +require_qdrant = pytest.mark.skipif(find_spec("qdrant_client") is None, reason="test requires qdrant-client") require_moto = pytest.mark.skipif(find_spec("moto") is None, reason="test requires moto") require_numpy1_on_windows = pytest.mark.skipif( version.parse(importlib.metadata.version("numpy")) >= version.parse("2.0.0") and sys.platform == "win32",