From d84f8e56caac0d3f78210d759efbae59661e30c2 Mon Sep 17 00:00:00 2001 From: Wangshu Pang Date: Fri, 16 Oct 2026 10:59:34 -0700 Subject: [PATCH] [Python] Add Qdrant search enrichment handler Add QdrantSearchEnrichmentHandler to ml.rag.enrichment so pipelines that ingest into Qdrant with QdrantWriteConfig can use Enrichment for the retrieval step. It supports dense, sparse and hybrid search; hybrid runs the dense and sparse searches as prefetches and fuses them server side with the Qdrant Query API (RRF by default, or DBSF). Each strategy takes a payload filter, limit, score_threshold and search_params, and hybrid search also takes a common filter and a threshold on the fused score. The handler reuses QdrantConnectionParameters from the ingestion module, its default vector names match QdrantWriteConfig, and each batch from BatchElements is sent as a single query_batch_points request. REST and gRPC quota and timeout errors are raised as UserCodeQuotaException and UserCodeTimeoutException so the Enrichment repeater retries them. join_fn stores results under metadata["enrichment_data"] and returns a copy of the item instead of mutating the input element. Builds on the earlier work in #35686. Fixes #40628 Co-authored-by: Anush Co-Authored-By: Claude Opus 5.5 --- CHANGES.md | 1 + .../ml/rag/enrichment/qdrant_search.py | 458 ++++++++++++++++++ .../rag/enrichment/qdrant_search_it_test.py | 210 ++++++++ .../ml/rag/enrichment/qdrant_search_test.py | 388 +++++++++++++++ 4 files changed, 1057 insertions(+) create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py create mode 100644 sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..57ebb5b747f5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -65,6 +65,7 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* Qdrant enrichment handler added (Python) ([#40628](https://github.com/apache/beam/issues/40628)). `QdrantSearchEnrichmentHandler` supports dense, sparse and hybrid (RRF/DBSF fusion) search with payload filters, and reads collections written by `QdrantWriteConfig`. * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py new file mode 100644 index 000000000000..85b510c105e3 --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search.py @@ -0,0 +1,458 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from dataclasses import dataclass +from dataclasses import field +from dataclasses import replace +from typing import Any +from typing import Optional +from typing import Union + +import grpc +import httpx +from qdrant_client import QdrantClient +from qdrant_client import models +from qdrant_client.common.client_exceptions import ResourceExhaustedResponse +from qdrant_client.http.exceptions import ResponseHandlingException +from qdrant_client.http.exceptions import UnexpectedResponse + +from apache_beam.io.requestresponse import UserCodeQuotaException +from apache_beam.io.requestresponse import UserCodeTimeoutException +from apache_beam.ml.rag.ingestion.qdrant import QdrantConnectionParameters +from apache_beam.ml.rag.types import EmbeddableItem +from apache_beam.transforms.enrichment import EnrichmentSourceHandler + +__all__ = [ + 'QdrantConnectionParameters', + 'QdrantDenseSearchParameters', + 'QdrantSparseSearchParameters', + 'QdrantHybridSearchParameters', + 'QdrantSearchParameters', + 'QdrantSearchResult', + 'QdrantSearchEnrichmentHandler', +] + + +@dataclass +class QdrantDenseSearchParameters: + """Parameters for dense vector similarity search. + + Args: + vector_name: Name of the dense vector in the Qdrant collection. Defaults to + "dense", matching the default `dense_embedding_key` of + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + limit: Maximum number of results to return per query. Must be positive. + Defaults to 3 search results. + filter: Optional payload filter applied to the search, e.g. + `models.Filter(must=[models.FieldCondition(key="source", + match=models.MatchValue(value="docs"))])`. + score_threshold: Optional minimal score; results scoring worse are dropped. + search_params: Optional `models.SearchParams` such as `hnsw_ef` or `exact`. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Enables forward compatibility. + """ + vector_name: str = "dense" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + search_params: Optional[models.SearchParams] = None + kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.vector_name: + raise ValueError("Vector name must be provided") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantSparseSearchParameters: + """Parameters for sparse vector (keyword) search. + + Args: + vector_name: Name of the sparse vector in the Qdrant collection. Defaults to + "sparse", matching the default `sparse_embedding_key` of + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + limit: Maximum number of results to return per query. Must be positive. + Defaults to 3 search results. + filter: Optional payload filter applied to the search. + score_threshold: Optional minimal score; results scoring worse are dropped. + search_params: Optional `models.SearchParams`. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Enables forward compatibility. + """ + vector_name: str = "sparse" + limit: int = 3 + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + search_params: Optional[models.SearchParams] = None + kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.vector_name: + raise ValueError("Vector name must be provided") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +@dataclass +class QdrantHybridSearchParameters: + """Parameters for hybrid (dense + sparse) search. + + The dense and sparse searches run as prefetches and their results are fused + server side with the Qdrant Query API. Configure `search_params` on `dense` + and `sparse` to control their respective prefetches; fusion does not run a + vector search. + + Args: + dense: Parameters for the dense search component. Its `limit` is the + number of candidates prefetched before fusion. + sparse: Parameters for the sparse search component. Its `limit` is the + number of candidates prefetched before fusion. + fusion: Fusion method used to combine both result lists. Defaults to + reciprocal rank fusion (`models.Fusion.RRF`). + limit: Maximum number of fused results to return per query. Defaults to 3 + search results. + filter: Optional common payload filter on the outer query. Qdrant also + applies it to the prefetches, together with their component filters. + score_threshold: Optional minimum fused score. Component thresholds apply + to the individual dense or sparse scores before fusion instead. + kwargs: Optional keyword arguments for additional `models.QueryRequest` + fields. Explicit non-None `filter` and `score_threshold` fields take + precedence over the corresponding entries in `kwargs`. + """ + dense: QdrantDenseSearchParameters + sparse: QdrantSparseSearchParameters + fusion: models.Fusion = models.Fusion.RRF + limit: int = 3 + kwargs: dict[str, Any] = field(default_factory=dict) + filter: Optional[models.Filter] = None + score_threshold: Optional[float] = None + + def __post_init__(self): + if not self.dense or not self.sparse: + raise ValueError( + "Dense and sparse search parameters must be provided for " + "hybrid search") + + if not self.fusion: + raise ValueError("Fusion must be provided for hybrid search") + + if self.limit <= 0: + raise ValueError(f"Search limit must be positive, got {self.limit}") + + +QdrantSearchStrategyType = Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters, + QdrantHybridSearchParameters] + + +@dataclass +class QdrantSearchParameters: + """Parameters configuring Qdrant search operations. + + Args: + collection_name: Name of the collection to search in. + search_strategy: Dense, sparse or hybrid search parameters. + with_payload: Whether to return the payload of each point. Either a bool + or a list of payload keys to return. Defaults to True. + consistency: Optional read consistency for the search, e.g. + `models.ReadConsistencyType.MAJORITY`. + timeout: Optional search timeout in seconds. If not specified, the client's + default timeout is used. + """ + collection_name: str + search_strategy: QdrantSearchStrategyType + with_payload: Union[bool, list[str]] = True + consistency: Optional[models.ReadConsistency] = None + timeout: Optional[int] = None + + def __post_init__(self): + if not self.collection_name: + raise ValueError("Collection name must be provided") + + if not self.search_strategy: + raise ValueError("Search strategy must be provided") + + +@dataclass +class QdrantSearchResult: + """Search result from Qdrant per embeddable_item. + + Args: + id: List of point IDs returned from the search. Can be either string (UUID) + or integer IDs. + score: List of scores for each returned point. Higher is more similar for + cosine and dot product, lower is more similar for Euclidean and Manhattan + distances. + payload: List of payload dictionaries, one for each returned point. Empty + if the point has no payload or `with_payload` is False. + """ + id: list[Union[str, int]] = field(default_factory=list) + score: list[float] = field(default_factory=list) + payload: list[dict[str, Any]] = field(default_factory=list) + + +InputT, OutputT = (Union[EmbeddableItem, list[EmbeddableItem]], + list[tuple[EmbeddableItem, dict[str, Any]]]) + + +class QdrantSearchEnrichmentHandler(EnrichmentSourceHandler[InputT, OutputT]): + """Enrichment handler for Qdrant vector database searches. + + This handler is designed to work with the + :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` transform. + It enables enriching data through dense, sparse, or hybrid searches against + Qdrant collections, such as ones written by + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + + The handler supports different search strategies: + * Dense search - For finding similar embeddings based on vector similarity + * Sparse search - For keyword-style retrieval with sparse embeddings + * Hybrid search - For fusing dense and sparse search results server side + + This handler queries Qdrant per element by default. To enable batching for + improved performance, set the `min_batch_size` and `max_batch_size` + parameters. These control the batching behavior in the + :class:`apache_beam.transforms.utils.BatchElements` transform. Each batch is + sent to Qdrant as a single batch query request. + + REST and gRPC quota and timeout errors are translated to Beam exceptions + so that `Enrichment` can retry them using its configured repeater. Other + client errors propagate unchanged. + """ + def __init__( + self, + connection_parameters: QdrantConnectionParameters, + search_parameters: QdrantSearchParameters, + *, + min_batch_size: int = 1, + max_batch_size: int = 1000, + **kwargs): + """ + Example Usage: + connection_parameters = QdrantConnectionParameters( + url="http://localhost:6333") + search_parameters = QdrantSearchParameters( + collection_name="my_collection", + search_strategy=QdrantDenseSearchParameters(limit=5)) + qdrant_handler = QdrantSearchEnrichmentHandler( + connection_parameters, + search_parameters, + min_batch_size=10, + max_batch_size=100) + + Args: + connection_parameters (QdrantConnectionParameters): Configuration for + connecting to Qdrant, shared with + :class:`~apache_beam.ml.rag.ingestion.qdrant.QdrantWriteConfig`. + search_parameters (QdrantSearchParameters): Configuration for search + operations, including collection name and search strategy. + min_batch_size (int): Minimum number of elements to batch together when + querying Qdrant. Default is 1 (no batching when max_batch_size is 1). + max_batch_size (int): Maximum number of elements to batch together. + Default is 1000. Higher values may improve throughput but increase + memory usage. + **kwargs: Additional keyword arguments for Qdrant Enrichment Handler. + """ + self._connection_parameters = connection_parameters + self._search_parameters = search_parameters + self._batching_kwargs = { + 'min_batch_size': min_batch_size, 'max_batch_size': max_batch_size + } + self.kwargs = kwargs + self._client = None + self.join_fn = join_fn + self.use_custom_types = True + + def __enter__(self): + """Enters the context manager and creates the Qdrant client. + + Returns: + Self, enabling use in 'with' statements. + """ + if not self._client: + params = self._connection_parameters + self._client = QdrantClient( + location=params.location, + url=params.url, + port=params.port, + grpc_port=params.grpc_port, + prefer_grpc=params.prefer_grpc, + https=params.https, + api_key=params.api_key, + prefix=params.prefix, + timeout=params.timeout, + host=params.host, + path=params.path, + check_compatibility=False, + **params.kwargs, + ) + return self + + def __call__( + self, + request: Union[EmbeddableItem, list[EmbeddableItem]], + *args, + **kwargs) -> list[tuple[EmbeddableItem, dict[str, Any]]]: + reqs = request if isinstance(request, list) else [request] + # Early return for empty requests to avoid unnecessary round trips. + if not reqs: + return [] + requests = [self._get_query_request(item) for item in reqs] + try: + responses = self._client.query_batch_points( + collection_name=self.collection_name, + requests=requests, + consistency=self._search_parameters.consistency, + timeout=self._search_parameters.timeout) + except ResourceExhaustedResponse as exc: + raise UserCodeQuotaException(str(exc)) from exc + except UnexpectedResponse as exc: + if exc.status_code == 429: + raise UserCodeQuotaException(str(exc)) from exc + if exc.status_code in (408, 504): + raise UserCodeTimeoutException(str(exc)) from exc + raise + except ResponseHandlingException as exc: + if isinstance(exc.source, httpx.TimeoutException): + raise UserCodeTimeoutException(str(exc)) from exc + raise + except grpc.RpcError as exc: + if exc.code() == grpc.StatusCode.RESOURCE_EXHAUSTED: + raise UserCodeQuotaException(str(exc)) from exc + if exc.code() == grpc.StatusCode.DEADLINE_EXCEEDED: + raise UserCodeTimeoutException(str(exc)) from exc + raise + return self._get_call_response(reqs, responses) + + def _get_query_request( + self, embeddable_item: EmbeddableItem) -> models.QueryRequest: + strategy = self.search_strategy + with_payload = self._search_parameters.with_payload + if isinstance(strategy, QdrantHybridSearchParameters): + query_kwargs = dict(strategy.kwargs) + if strategy.filter is not None: + query_kwargs['filter'] = strategy.filter + if strategy.score_threshold is not None: + query_kwargs['score_threshold'] = strategy.score_threshold + return models.QueryRequest( + prefetch=[ + self._get_prefetch( + strategy.dense, self._get_dense_query(embeddable_item)), + self._get_prefetch( + strategy.sparse, self._get_sparse_query(embeddable_item)), + ], + query=models.FusionQuery(fusion=strategy.fusion), + limit=strategy.limit, + with_payload=with_payload, + **query_kwargs) + elif isinstance(strategy, QdrantDenseSearchParameters): + return self._get_single_query_request( + strategy, self._get_dense_query(embeddable_item), with_payload) + elif isinstance(strategy, QdrantSparseSearchParameters): + return self._get_single_query_request( + strategy, self._get_sparse_query(embeddable_item), with_payload) + else: + raise ValueError(f"Not supported search strategy yet: {strategy}") + + def _get_single_query_request( + self, + strategy: Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters], + query: Union[list[float], models.SparseVector], + with_payload: Union[bool, list[str]]) -> models.QueryRequest: + return models.QueryRequest( + query=query, + using=strategy.vector_name, + filter=strategy.filter, + score_threshold=strategy.score_threshold, + params=strategy.search_params, + limit=strategy.limit, + with_payload=with_payload, + **strategy.kwargs) + + def _get_prefetch( + self, + strategy: Union[QdrantDenseSearchParameters, + QdrantSparseSearchParameters], + query: Union[list[float], models.SparseVector]) -> models.Prefetch: + return models.Prefetch( + query=query, + using=strategy.vector_name, + filter=strategy.filter, + score_threshold=strategy.score_threshold, + params=strategy.search_params, + limit=strategy.limit, + **strategy.kwargs) + + def _get_dense_query(self, embeddable_item: EmbeddableItem) -> list[float]: + if not embeddable_item.dense_embedding: + raise ValueError( + f"Item {embeddable_item.id} missing dense embedding required for" + " dense search") + return embeddable_item.dense_embedding + + def _get_sparse_query( + self, embeddable_item: EmbeddableItem) -> models.SparseVector: + if not embeddable_item.sparse_embedding: + raise ValueError( + f"Item {embeddable_item.id} missing sparse embedding required for" + " sparse search") + indices, values = embeddable_item.sparse_embedding + return models.SparseVector(indices=indices, values=values) + + def _get_call_response( + self, + embeddable_items: list[EmbeddableItem], + responses: list[models.QueryResponse]): + response = [] + for embeddable_item, query_response in zip(embeddable_items, responses): + result = QdrantSearchResult() + for point in query_response.points: + result.id.append(point.id) + result.score.append(point.score) + result.payload.append(point.payload or {}) + response.append((embeddable_item, result.__dict__)) + return response + + @property + def collection_name(self): + """Getter method for collection_name property""" + return self._search_parameters.collection_name + + @property + def search_strategy(self): + """Getter method for search_strategy property""" + return self._search_parameters.search_strategy + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._client: + try: + self._client.close() + finally: + self._client = None + + def batch_elements_kwargs(self) -> dict[str, int]: + """Returns kwargs for beam.BatchElements.""" + return self._batching_kwargs + + +def join_fn(left: EmbeddableItem, right: dict[str, Any]) -> EmbeddableItem: + return replace(left, metadata={**left.metadata, 'enrichment_data': right}) diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py new file mode 100644 index 000000000000..2d423ec2c5ab --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_it_test.py @@ -0,0 +1,210 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest + +import pytest + +import apache_beam as beam +from apache_beam.ml.rag.types import Content +from apache_beam.ml.rag.types import EmbeddableItem +from apache_beam.ml.rag.types import Embedding +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +# pylint: disable=ungrouped-imports +try: + from qdrant_client import models + + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantConnectionParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantDenseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantHybridSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchEnrichmentHandler + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSparseSearchParameters + from apache_beam.ml.rag.ingestion.qdrant import QdrantWriteConfig + from apache_beam.transforms.enrichment import Enrichment + + QDRANT_AVAILABLE = True +except ImportError: + QDRANT_AVAILABLE = False + +try: + from testcontainers.qdrant import QdrantContainer +except ImportError: + QdrantContainer = None +# pylint: enable=ungrouped-imports + +TEST_CORPUS = [ + EmbeddableItem( + id="1", + content=Content(text="Beam runs batch and streaming pipelines."), + metadata={"source": "beam"}, + embedding=Embedding( + dense_embedding=[1.0, 0.0], sparse_embedding=([0, 1], [1.0, 0.5])), + ), + EmbeddableItem( + id="2", + content=Content(text="Qdrant is a vector database."), + metadata={"source": "qdrant"}, + embedding=Embedding( + dense_embedding=[0.0, 1.0], sparse_embedding=([2, 3], [1.0, 0.5])), + ), + EmbeddableItem( + id="3", + content=Content(text="Enrichment joins external data."), + metadata={"source": "beam"}, + embedding=Embedding( + dense_embedding=[-1.0, 0.0], sparse_embedding=([4, 5], [1.0, 0.5])), + ), +] + + +def _query(id, dense=None, sparse=None): + return EmbeddableItem( + id=id, + content=Content(text=f"query {id}"), + embedding=Embedding(dense_embedding=dense, sparse_embedding=sparse)) + + +def _retrieved(item): + """Extracts (query id, retrieved ids, retrieved contents) for assertions.""" + data = item.metadata["enrichment_data"] + return ( + item.id, + data["id"], [payload.get("content") for payload in data["payload"]]) + + +@unittest.skipIf(not QDRANT_AVAILABLE, "qdrant dependencies not installed.") +@unittest.skipIf(not QdrantContainer, "qdrant test_container not installed.") +@pytest.mark.require_docker_in_docker +class TestQdrantSearchEnrichment(unittest.TestCase): + """Writes TEST_CORPUS with QdrantWriteConfig, then retrieves it.""" + _collection_name = "test_qdrant_search" + + @classmethod + def setUpClass(cls): + cls._container = QdrantContainer() + cls._container.start() + cls.addClassCleanup(cls._container.stop) + cls._connection_params = QdrantConnectionParameters( + host=cls._container.get_container_host_ip(), + port=int(cls._container.get_exposed_port(6333))) + + client = cls._container.get_client() + client.create_collection( + collection_name=cls._collection_name, + vectors_config={ + "dense": models.VectorParams( + size=2, distance=models.Distance.COSINE) + }, + sparse_vectors_config={"sparse": models.SparseVectorParams()}) + client.close() + + write_config = QdrantWriteConfig( + connection_params=cls._connection_params, + collection_name=cls._collection_name, + content_payload_key="content") + with TestPipeline() as p: + _ = (p | beam.Create(TEST_CORPUS) | write_config.create_write_transform()) + + def _run_enrichment(self, search_strategy, queries, expected, **kwargs): + handler = QdrantSearchEnrichmentHandler( + self._connection_params, + QdrantSearchParameters( + collection_name=self._collection_name, + search_strategy=search_strategy), + **kwargs) + with TestPipeline() as p: + result = ( + p + | beam.Create(queries) + | Enrichment(handler) + | beam.Map(_retrieved)) + assert_that(result, equal_to(expected)) + + def test_dense_search(self): + self._run_enrichment( + QdrantDenseSearchParameters(limit=1), + [_query("a", dense=[0.9, 0.1]), _query("b", dense=[0.1, 0.9])], + [ + ("a", [1], ["Beam runs batch and streaming pipelines."]), + ("b", [2], ["Qdrant is a vector database."]), + ]) + + def test_dense_search_with_filter(self): + self._run_enrichment( + QdrantDenseSearchParameters( + limit=3, + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="beam")) + ])), [_query("a", dense=[0.1, 0.9])], + [( + "a", [1, 3], + [ + "Beam runs batch and streaming pipelines.", + "Enrichment joins external data." + ])]) + + def test_sparse_search(self): + self._run_enrichment( + QdrantSparseSearchParameters(limit=3), + [_query("a", sparse=([4], [1.0]))], + [("a", [3], ["Enrichment joins external data."])]) + + def test_hybrid_search(self): + self._run_enrichment( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + limit=1), + [_query("a", dense=[0.0, 1.0], sparse=([2, 3], [1.0, 1.0]))], + [("a", [2], ["Qdrant is a vector database."])]) + + def test_batched_search(self): + queries = [_query(str(i), dense=[1.0, 0.0]) for i in range(10)] + self._run_enrichment( + QdrantDenseSearchParameters(limit=1), + queries, + [(q.id, [1], ["Beam runs batch and streaming pipelines."]) + for q in queries], + min_batch_size=5, + max_batch_size=5) + + def test_hybrid_search_with_filter_and_fused_score_threshold(self): + self._run_enrichment( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="beam")) + ]), + score_threshold=0.7), + [_query("a", dense=[0.1, 0.9], sparse=([4], [1.0]))], + [("a", [3], ["Enrichment joins external data."])]) + + def test_empty_input(self): + self._run_enrichment(QdrantDenseSearchParameters(), [], []) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py new file mode 100644 index 000000000000..0c3a6a706c88 --- /dev/null +++ b/sdks/python/apache_beam/ml/rag/enrichment/qdrant_search_test.py @@ -0,0 +1,388 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +from unittest.mock import patch + +import grpc +from parameterized import parameterized + +try: + import httpx + from qdrant_client import models + from qdrant_client.common.client_exceptions import ResourceExhaustedResponse + from qdrant_client.http.exceptions import ResponseHandlingException + from qdrant_client.http.exceptions import UnexpectedResponse + from qdrant_client.http.models import QueryResponse + + from apache_beam.io.requestresponse import ExponentialBackOffRepeater + from apache_beam.io.requestresponse import UserCodeQuotaException + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantConnectionParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantDenseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantHybridSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchEnrichmentHandler + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import QdrantSparseSearchParameters + from apache_beam.ml.rag.enrichment.qdrant_search import join_fn + from apache_beam.ml.rag.types import Content + from apache_beam.ml.rag.types import EmbeddableItem + from apache_beam.ml.rag.types import Embedding +except ImportError as e: + raise unittest.SkipTest(f'Qdrant dependencies not installed: {str(e)}') + +_COLLECTION = "test_collection" + +# Points stored in the in-memory collection. Dense vectors use cosine +# distance, so a query of [1.0, 0.0] ranks point 1 first, then 2, then 3. +_POINTS = [ + models.PointStruct( + id=1, + vector={ + "dense": [1.0, 0.0], + "sparse": models.SparseVector(indices=[0, 1], values=[1.0, 0.5]), + }, + payload={ + "content": "doc one", "source": "a" + }, + ), + models.PointStruct( + id=2, + vector={ + "dense": [0.8, 0.6], + "sparse": models.SparseVector(indices=[1, 2], values=[1.0, 0.5]), + }, + payload={ + "content": "doc two", "source": "b" + }, + ), + models.PointStruct( + id=3, + vector={ + "dense": [-1.0, 0.0], + "sparse": models.SparseVector(indices=[2, 3], values=[1.0, 0.5]), + }, + payload={ + "content": "doc three", "source": "a" + }, + ), +] + + +def _item(id, dense=None, sparse=None): + return EmbeddableItem( + id=id, + content=Content(text=f"query {id}"), + embedding=Embedding(dense_embedding=dense, sparse_embedding=sparse)) + + +class _RpcError(grpc.RpcError): + def __init__(self, status_code): + self._status_code = status_code + + def code(self): + return self._status_code + + +class TestQdrantSearchParameters(unittest.TestCase): + """Unit tests for validation of Qdrant search parameters.""" + @parameterized.expand([ + (lambda: QdrantDenseSearchParameters(limit=0), "positive"), + (lambda: QdrantDenseSearchParameters(vector_name=""), "Vector name"), + (lambda: QdrantSparseSearchParameters(limit=-1), "positive"), + (lambda: QdrantSparseSearchParameters(vector_name=""), "Vector name"), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), sparse= + QdrantSparseSearchParameters(), limit=0), + "positive"), + ( + lambda: QdrantHybridSearchParameters( + dense=None, sparse=QdrantSparseSearchParameters()), + "Dense and sparse"), + ( + lambda: QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), sparse= + QdrantSparseSearchParameters(), fusion=None), + "Fusion"), + ( + lambda: QdrantSearchParameters( + collection_name="", search_strategy=QdrantDenseSearchParameters() + ), + "Collection name"), + ( + lambda: QdrantSearchParameters( + collection_name=_COLLECTION, search_strategy=None), + "Search strategy"), + ]) + def test_invalid_parameters(self, create, expected_message): + with self.assertRaises(ValueError) as context: + create() + self.assertIn(expected_message, str(context.exception)) + + def test_batch_elements_kwargs(self): + handler = QdrantSearchEnrichmentHandler( + QdrantConnectionParameters.in_memory(), + QdrantSearchParameters( + collection_name=_COLLECTION, + search_strategy=QdrantDenseSearchParameters()), + min_batch_size=5, + max_batch_size=50) + self.assertEqual( + handler.batch_elements_kwargs(), { + 'min_batch_size': 5, 'max_batch_size': 50 + }) + + +class TestQdrantSearchEnrichmentHandler(unittest.TestCase): + """Runs the handler against an in-memory Qdrant collection.""" + def _create_handler(self, search_strategy, **search_kwargs): + handler = QdrantSearchEnrichmentHandler( + QdrantConnectionParameters.in_memory(), + QdrantSearchParameters( + collection_name=_COLLECTION, + search_strategy=search_strategy, + **search_kwargs)) + handler.__enter__() + self.addCleanup(handler.__exit__, None, None, None) + # The in-memory client is created on __enter__, so populate it here. + client = handler._client + client.create_collection( + collection_name=_COLLECTION, + vectors_config={ + "dense": models.VectorParams( + size=2, distance=models.Distance.COSINE) + }, + sparse_vectors_config={"sparse": models.SparseVectorParams()}) + client.upsert(collection_name=_COLLECTION, points=_POINTS) + return handler + + def test_dense_search(self): + handler = self._create_handler(QdrantDenseSearchParameters(limit=2)) + item = _item("q", dense=[1.0, 0.0]) + + [(returned_item, result)] = handler(item) + + self.assertIs(returned_item, item) + self.assertEqual(result["id"], [1, 2]) + self.assertAlmostEqual(result["score"][0], 1.0, places=5) + self.assertAlmostEqual(result["score"][1], 0.8, places=5) + self.assertEqual( + result["payload"][0], { + "content": "doc one", "source": "a" + }) + + def test_dense_search_with_filter_and_score_threshold(self): + handler = self._create_handler( + QdrantDenseSearchParameters( + limit=3, + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="a")) + ]), + score_threshold=0.0)) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + # Point 2 is filtered out by payload and point 3 by score (-1.0). + self.assertEqual(result["id"], [1]) + + def test_sparse_search(self): + handler = self._create_handler(QdrantSparseSearchParameters(limit=3)) + + [(_, result)] = handler(_item("q", sparse=([2, 3], [1.0, 1.0]))) + + # Dot products: point 3 = 1.5, point 2 = 0.5, point 1 shares no index. + self.assertEqual(result["id"], [3, 2]) + self.assertEqual(result["score"], [1.5, 0.5]) + + @parameterized.expand([(models.Fusion.RRF, ), (models.Fusion.DBSF, )]) + def test_hybrid_search(self, fusion): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + fusion=fusion, + limit=2)) + + [(_, result) + ] = handler(_item("q", dense=[0.8, 0.6], sparse=([1, 2], [1.0, 1.0]))) + + # Point 2 ranks first in both prefetches, so it wins after fusion. + self.assertEqual(len(result["id"]), 2) + self.assertEqual(result["id"][0], 2) + self.assertEqual(result["payload"][0]["content"], "doc two") + + def test_hybrid_common_filter_applies_before_prefetch_limit(self): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=1), + sparse=QdrantSparseSearchParameters(limit=1), + filter=models.Filter( + must=[ + models.FieldCondition( + key="source", match=models.MatchValue(value="a")) + ]))) + + [(_, result)] = handler(_item("q", dense=[0.8, 0.6], sparse=([1], [1.0]))) + + # Point 2 would win both prefetches without the common filter. + self.assertEqual(result["id"], [1]) + + def test_hybrid_fused_score_threshold(self): + handler = self._create_handler( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(limit=3), + sparse=QdrantSparseSearchParameters(limit=3), + score_threshold=0.8)) + + [(_, result) + ] = handler(_item("q", dense=[0.8, 0.6], sparse=([1, 2], [1.0, 1.0]))) + + self.assertEqual(result["id"], [2]) + self.assertEqual(result["score"], [1.0]) + + @parameterized.expand([ + (ResourceExhaustedResponse("rate limited", 1), ), + (UnexpectedResponse(429, "rate limited", b"", httpx.Headers()), ), + (UnexpectedResponse(408, "timeout", b"", httpx.Headers()), ), + (UnexpectedResponse(504, "timeout", b"", httpx.Headers()), ), + (ResponseHandlingException(httpx.ReadTimeout("timeout")), ), + (_RpcError(grpc.StatusCode.RESOURCE_EXHAUSTED), ), + (_RpcError(grpc.StatusCode.DEADLINE_EXCEEDED), ), + ]) + def test_default_repeater_recovers_from_quota_and_timeout(self, error): + handler = self._create_handler(QdrantDenseSearchParameters(limit=1)) + item = _item("q", dense=[1.0, 0.0]) + response = QueryResponse(points=[]) + with patch.object( + handler._client, 'query_batch_points', + side_effect=[error, [response]]) as query, patch( + 'apache_beam.utils.retry.time.sleep'): + result = ExponentialBackOffRepeater().repeat(handler, [item], 30) + + self.assertEqual(result, [(item, {"id": [], "score": [], "payload": []})]) + self.assertEqual(query.call_count, 2) + + @parameterized.expand([ + (UnexpectedResponse(400, "bad request", b"", httpx.Headers()), ), + (UnexpectedResponse(401, "unauthorized", b"", httpx.Headers()), ), + (UnexpectedResponse(404, "not found", b"", httpx.Headers()), ), + (ResponseHandlingException(ValueError("invalid response")), ), + (_RpcError(grpc.StatusCode.INVALID_ARGUMENT), ), + (ValueError("invalid query"), ), + ]) + def test_default_repeater_does_not_retry_permanent_errors(self, error): + handler = self._create_handler(QdrantDenseSearchParameters()) + with patch.object(handler._client, 'query_batch_points', + side_effect=error) as query: + with self.assertRaises(type(error)) as context: + ExponentialBackOffRepeater().repeat( + handler, [_item("q", dense=[1.0, 0.0])], 30) + + self.assertIs(context.exception, error) + self.assertEqual(query.call_count, 1) + + def test_default_repeater_exhausts_retries(self): + handler = self._create_handler(QdrantDenseSearchParameters()) + error = ResourceExhaustedResponse("rate limited", 1) + with patch.object(handler._client, 'query_batch_points', + side_effect=error) as query, \ + patch('apache_beam.utils.retry.time.sleep'): + with self.assertRaises(UserCodeQuotaException) as context: + ExponentialBackOffRepeater().repeat( + handler, [_item("q", dense=[1.0, 0.0])], 30) + + self.assertIs(context.exception.__cause__, error) + self.assertEqual(query.call_count, 3) + + def test_batched_request_returns_result_per_item(self): + handler = self._create_handler(QdrantDenseSearchParameters(limit=1)) + items = [ + _item("q1", dense=[1.0, 0.0]), + _item("q2", dense=[-1.0, 0.0]), + ] + + response = handler(items) + + self.assertEqual([item for item, _ in response], items) + self.assertEqual([result["id"] for _, result in response], [[1], [3]]) + + def test_with_payload_keys(self): + handler = self._create_handler( + QdrantDenseSearchParameters(limit=1), with_payload=["content"]) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + self.assertEqual(result["payload"], [{"content": "doc one"}]) + + def test_without_payload(self): + handler = self._create_handler( + QdrantDenseSearchParameters(limit=1), with_payload=False) + + [(_, result)] = handler(_item("q", dense=[1.0, 0.0])) + + self.assertEqual(result["id"], [1]) + self.assertEqual(result["payload"], [{}]) + + def test_empty_request(self): + handler = self._create_handler(QdrantDenseSearchParameters()) + self.assertEqual(handler([]), []) + + @parameterized.expand([ + (QdrantDenseSearchParameters(), _item("q", sparse=([0], [1.0])), "dense"), + (QdrantSparseSearchParameters(), _item("q", dense=[1.0, 0.0]), "sparse"), + ( + QdrantHybridSearchParameters( + dense=QdrantDenseSearchParameters(), + sparse=QdrantSparseSearchParameters()), + _item("q", dense=[1.0, 0.0]), + "sparse"), + ]) + def test_missing_embedding(self, search_strategy, item, embedding_type): + handler = self._create_handler(search_strategy) + with self.assertRaises(ValueError) as context: + handler(item) + self.assertIn( + f"Item q missing {embedding_type} embedding", str(context.exception)) + + def test_unsupported_search_strategy(self): + handler = self._create_handler({"vector_name": "dense"}) + with self.assertRaises(ValueError) as context: + handler(_item("q", dense=[1.0, 0.0])) + self.assertIn("Not supported search strategy", str(context.exception)) + + def test_join_fn(self): + item = _item("q", dense=[1.0, 0.0]) + item.metadata = {"source": "query", "enrichment_data": {"previous": True}} + original_metadata = dict(item.metadata) + result = {"id": [1], "score": [1.0], "payload": [{"content": "doc one"}]} + + joined = join_fn(item, result) + + self.assertEqual(joined.metadata["enrichment_data"], result) + self.assertEqual(joined.metadata["source"], "query") + self.assertEqual(item.metadata, original_metadata) + self.assertIsNot(joined, item) + self.assertIsNot(joined.metadata, item.metadata) + self.assertEqual(joined.id, item.id) + self.assertEqual(joined.content, item.content) + self.assertEqual(joined.embedding, item.embedding) + + +if __name__ == '__main__': + unittest.main()