pangwangshu opened a new pull request, #40629: URL: https://github.com/apache/beam/pull/40629
Adds `QdrantSearchEnrichmentHandler` to `apache_beam/ml/rag/enrichment`, so a pipeline that ingests into Qdrant with `QdrantWriteConfig` can use `Enrichment` for the retrieval step. Until now the enrichment handlers covered only BigQuery and Milvus. This builds on #35686 by @anush008, which was closed as stale before review. Anush agreed to my picking it up ([comment](https://github.com/apache/beam/pull/35686#issuecomment-6207317367)) and is credited as a co-author on the commit. **Design** The handler follows `MilvusSearchEnrichmentHandler`, so the two read the same way: | Milvus | Qdrant | | --- | --- | | `VectorSearchParameters` | `QdrantDenseSearchParameters` | | `KeywordSearchParameters` | `QdrantSparseSearchParameters` | | `HybridSearchParameters` + ranker | `QdrantHybridSearchParameters` + `fusion` (`RRF` default, or `DBSF`) | | `MilvusSearchParameters` | `QdrantSearchParameters` | | `MilvusSearchResult` (`id`, `distance`, `fields`) | `QdrantSearchResult` (`id`, `score`, `payload`) | - Hybrid search sends the dense and sparse queries as prefetches and fuses the results server side through the Qdrant Query API. - Each strategy accepts a payload `filter`, `limit`, `score_threshold`, `search_params`, and `kwargs` for other `QueryRequest` fields. `with_payload` can be a bool or a list of payload keys. - `QdrantConnectionParameters` is reused from `ingestion/qdrant.py` and re-exported. The default vector names (`dense` / `sparse`) match `QdrantWriteConfig`, so a collection written by Beam can be searched without extra configuration. - Each batch from `BatchElements` is sent as a single `query_batch_points` request rather than one request per element. - `join_fn`, the `__call__` return shape, empty-input handling and `batch_elements_kwargs` match the Milvus handler. - Sparse search requires a sparse embedding. Unlike Milvus, Qdrant has no built-in BM25 over raw text. Compared with #35686: the handler uses `EmbeddableItem` instead of the old `Chunk` type, reuses the existing connection parameters instead of duplicating them, and batches queries. No `setup.py` change is needed, because `qdrant-client` is already part of the `qdrant` extra. A website example page, like the one for Milvus enrichment, can follow in a separate PR. **Testing** - `qdrant_search_test.py`: 22 unit tests covering parameter validation, plus dense, sparse, hybrid, filter, `score_threshold`, `with_payload`, batching, empty input and missing-embedding errors. They run against Qdrant's in-memory mode, so they need no Docker. All pass locally. - `qdrant_search_it_test.py`: 6 integration tests. Each writes a corpus with `QdrantWriteConfig(content_payload_key="content")` to a testcontainers Qdrant instance, then reads it back through `Enrichment` (dense, filtered dense, sparse, hybrid, batched, empty input). All pass locally on DirectRunner. - yapf, pylint and ruff pass on the new files. - `CHANGES.md` updated. **AI assistance disclosure** Per the [ASF generative tooling guidance](https://www.apache.org/legal/generative-tooling.html): this change was written with the help of Claude Code (Claude Opus 5.5). I have reviewed the code and tests and ran them locally, and to my knowledge the change contains no third-party copyrighted material. Generated-by: Claude Code (Claude Opus 5.5) Fixes #40628 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
