Skip to content

[Feature Request]: Add a Qdrant search enrichment handler for ml.rag #40628

Open

Description

@pangwangshu

What would you like to happen?

ml.rag can write to Qdrant (QdrantWriteConfig, #38142; chunk content via content_payload_key, #40502), but there is no way to retrieve from it. The enrichment handlers in apache_beam/ml/rag/enrichment cover only BigQuery and Milvus, so a pipeline that ingests into Qdrant cannot use Enrichment for the retrieval step.

An earlier attempt, #35686, was closed as stale before it was reviewed. @anush008 has agreed to my picking it up in a new PR (#35686 (comment)).

Proposal

Add QdrantSearchEnrichmentHandler in apache_beam/ml/rag/enrichment/qdrant_search.py, modeled on MilvusSearchEnrichmentHandler:

  • Search strategies:
    • QdrantDenseSearchParameters: dense vector search.
    • QdrantSparseSearchParameters: sparse vector search.
    • QdrantHybridSearchParameters: dense and sparse prefetches fused server side through the Qdrant Query API (RRF by default, or DBSF).
  • Each strategy accepts a payload filter, limit, score_threshold and search_params. with_payload can be a bool or a list of payload keys.
  • Reuses QdrantConnectionParameters from ingestion/qdrant.py. The default vector names (dense / sparse) match QdrantWriteConfig, so collections written by Beam can be searched without extra configuration.
  • Sends each batch from BatchElements as one query_batch_points request rather than one request per element.
  • Each element gets {"id": [...], "score": [...], "payload": [...]}, attached by the same join_fn (metadata["enrichment_data"]) as the Milvus handler.
  • Tests: unit tests that run against Qdrant's in-memory mode (no Docker needed), and a testcontainers integration test that writes with QdrantWriteConfig and reads back through Enrichment.
  • qdrant-client is already a dependency (qdrant extra), so no new dependencies or requirements regeneration are needed.

A website example page (like the one for Milvus enrichment) can follow in a separate PR.

I have an implementation ready and plan to send a PR.

Issue Priority

Priority: 2 (default / most feature requests should be filed as P2)

Issue Components

  • Component: Python SDK
  • Component: Java SDK
  • Component: Go SDK
  • Component: Typescript SDK
  • Component: IO connector
  • Component: Beam YAML
  • Component: Beam examples
  • Component: Beam playground
  • Component: Beam katas
  • Component: Website
  • Component: Infrastructure
  • Component: Spark Runner
  • Component: Flink Runner
  • Component: Prism Runner
  • Component: Twister2 Runner
  • Component: Hazelcast Jet Runner
  • Component: Google Cloud Dataflow Runner

Activity

  1. pangwangshu commented on Oct 16, 2026

    @pangwangshu
    ContributorAuthor

    .take-issue

  2. added a commit that references this issue on Oct 16, 2026
    d84f8e5
  3. pangwangshu commented on Oct 16, 2026

    @pangwangshu
    ContributorAuthor

    I've opened #40629 for this.

    It adds QdrantSearchEnrichmentHandler with dense, sparse and hybrid (RRF/DBSF) search. It reuses QdrantConnectionParameters from the ingestion module, so collections written by QdrantWriteConfig can be searched without extra configuration. Unit tests run against Qdrant's in-memory mode, and the integration tests write with QdrantWriteConfig and read back through Enrichment.

    @anush008, this builds on your work in #35686, and the commit credits you as a co-author. Thanks again for starting it!

    @derrickaw, since you looked at #35686, could you take a look at #40629 or point it to the right reviewer when you have a chance? Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions