Skip to content

Add QdrantSearchOperator for vector similarity search - #69673

Closed
YAshhh29 wants to merge 8 commits into
apache:mainfrom
YAshhh29:feature/qdrant-search-operator
Closed

YAshhh29 wants to merge 8 commits into
apache:mainfrom
YAshhh29:feature/qdrant-search-operator

Conversation

@YAshhh29

@YAshhh29 YAshhh29 commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

The Qdrant provider today lets users write vectors into a collection
(QdrantIngestOperator) but has no operator for the other half of a RAG
pipeline: reading them back out. Users who want to run a similarity search
from a DAG have to reach into hook.conn.query_points directly and remember
to convert the returned pydantic ScoredPoint objects into plain dicts so
Airflow can serialize them to XCom -- a footgun that shows up as a cryptic
serialization error at runtime.

This PR adds QdrantSearchOperator, a first-class task that closes that
gap. Every other vector-DB provider (Pinecone, Weaviate) is missing the
same operator; Qdrant is the leanest of the three (its hook doesn't even
have a search method today), so it's the cleanest place to start.

How I found and verified this gap

This isn't tied to an existing issue -- I discovered it by auditing the
operator surface of every AI/ML provider in Airflow (openai, cohere,
pinecone, weaviate, qdrant, pgvector), the same audit approach behind
#69408 and #69534. For each provider I compared what the hook/client can
do to what's actually exposed as operators.

The pattern that jumped out: every vector-DB provider is missing a
search operator
. Users can ingest with a proper operator but must fall
back to raw hook calls to query. That's the retrieval half of RAG living
outside the Airflow abstraction.

Before writing a line of code I confirmed:

  1. No competing work in flight. GitHub search returned 0 open PRs and
    0 open issues mentioning "qdrant search" -- greenfield, no one else
    was building this.
  2. The upstream API is stable and modern. qdrant-client 1.18.0
    (the provider pins >=1.17.1) exposes query_points with every one
    of the 9 named parameters this operator forwards; the older search()
    method is deprecated and slated for removal.
  3. The response contract is what I assumed. QueryResponse.points
    is List[ScoredPoint], and ScoredPoint.model_dump() produces the
    id/score/payload/vector/version/shard_key/order_value dict shape
    the operator promises callers.
  4. The provider's registry auto-discovers by module, not by class.
    python-modules in provider.yaml covers any class in
    operators/qdrant.py, so adding one needs zero registry edits.

Only then did I write the code, in the small incremental steps you can
see in the six commits (hook -> hook tests -> operator -> operator tests
-> example DAG -> docs).

Design decisions

  • A hook method + a thin operator, not just an operator. A new
    QdrantHook.search() wraps QdrantClient.query_points and converts
    each returned ScoredPoint to a plain dict via model_dump(). The
    operator is a ~10-line delegate on top. This mirrors the operator/hook
    split every other provider uses -- and gives tests a clean seam to
    mock at.
  • XCom-safe by construction. The hook returns list[dict[str, Any]]
    (id, score, payload, and optionally vector), so results land in XCom
    without any user-side workaround.
  • Uses query_points, not the deprecated search(). The search()
    API in qdrant-client is scheduled for removal in a future major;
    query_points is the modern surface (also supports named/sparse
    vectors, hybrid search, etc.). A regression test asserts we never
    fall back to the deprecated method.
  • **kwargs passthrough. Forwards any query_points parameter we
    don't enumerate (using, prefetch, lookup_from, ...) so the hook
    stays forward-compatible with hybrid search and named vectors without
    a follow-up PR.

What changes

  • providers/qdrant/src/airflow/providers/qdrant/hooks/qdrant.py
    • New QdrantHook.search(...) method wrapping query_points, returning
      list[dict] via ScoredPoint.model_dump().
  • providers/qdrant/src/airflow/providers/qdrant/operators/qdrant.py
    • New QdrantSearchOperator class alongside the existing
      QdrantIngestOperator. template_fields include collection_name,
      query, query_filter, limit so a RAG DAG can XCom-pull a query
      vector from an upstream embedding task.
  • providers/qdrant/tests/unit/qdrant/hooks/test_qdrant.py
    • Three tests: return type is list[dict] via model_dump; uses
      query_points (not deprecated search) with all named args
      forwarded; extra **kwargs also forwarded.
  • providers/qdrant/tests/unit/qdrant/operators/test_qdrant.py
    • Five tests: execute returns the hook result; defaults forward as
      expected; every optional arg reaches the hook; template_fields
      cover the runtime parameters; default conn_id matches the hook's.
  • providers/qdrant/tests/system/qdrant/example_dag_qdrant.py
    • Adds a QdrantSearchOperator task downstream of the existing
      ingest task with # [START/END] howto_operator_qdrant_search
      markers.
  • providers/qdrant/docs/operators/qdrant.rst
    • How-to section with the matching .. _howto/operator:QdrantSearchOperator:
      anchor and an .. exampleinclude:: pulling the DAG snippet.

No provider.yaml / get_provider_info.py changes needed: the registry
lists python-modules, not classes, so a new class in an existing module is
picked up automatically. No changelog edit either -- provider changelogs
are regenerated from git log by the release manager per AGENTS.md.

Testing

  • All 8 unit tests pass locally (3 hook + 5 operator), verified via
    a standalone harness that runs the real hook/operator code with mocked
    Qdrant client + a BaseHook/BaseOperator shim (full Airflow can't
    run on Windows).
  • API contract verified against qdrant-client 1.18.0: query_points
    accepts every one of the 9 named parameters we forward, and
    QueryResponse.points is a List[ScoredPoint] with the expected
    model_dump() shape (id, score, payload, vector, ...).
  • Regression check: QdrantIngestOperator still constructs and
    behaves identically -- we only added to the module, no existing code
    was touched.
  • Full-provider ruff check + ruff format --check: 26 files clean.
  • Self-reviewed against every rule in .github/instructions/code-review.instructions.md
    -- no red flags (no time.time, no assert in prod, no new
    AirflowException, no British spellings, no missing tests).

Was generative AI tooling used to co-author this PR?
  • Yes -- GitHub Copilot (Claude Opus 4.6)

Generated-by: GitHub Copilot (Claude Opus 4.6) following the guidelines

@YAshhh29

Copy link
Copy Markdown
Contributor Author

Rebased onto latest main; all checks are green. This adds a first-class QdrantSearchOperator (the retrieval half of a RAG pipeline) alongside the existing ingest operator - a small, self-contained addition with full unit coverage.
It already carries the ready for maintainer review label; would appreciate a review from a maintainer whenever someone has a moment. Thanks!

apache#73286 made every connection id an operator accepts a template field, and added
the check-conn-id-templated hook to enforce it; QdrantIngestOperator gained
conn_id in that change. QdrantSearchOperator accepted conn_id without
templating it, so the hook fails on this branch once it is rebased onto main.
Add it last, matching QdrantIngestOperator.
The test documents which fields users are expected to template, and conn_id is
now one of them. Without it in the expected set, dropping conn_id from
template_fields would be caught only by the static check, not by the unit tests.
@YAshhh29
YAshhh29 force-pushed the feature/qdrant-search-operator branch from 6d6d8c4 to c05a38f Compare September 25, 2026 16:19
@YAshhh29

Copy link
Copy Markdown
Contributor Author

Rebased onto main. #73286 has since made every connection id an operator accepts a template field, and added a check-conn-id-templated hook to enforce it — this PR's QdrantSearchOperator would have failed it, so conn_id is now templated, matching QdrantIngestOperator, and the template-fields test covers it.

Re-checked after the rebase: the qdrant unit tests pass against qdrant-client 1.19.1, check-conn-id-templated, validate-operators-init and the template-fields check pass, and ruff and mypy are clean. Would appreciate a review when someone has a moment.

@potiuk potiuk added the closed because of open PR limit Closed as a one-time step of introducing the open pull request limit label Sep 25, 2026
@potiuk

potiuk commented Sep 25, 2026

Copy link
Copy Markdown
Member

Hello @YAshhh29 - thank you for your contributions to Apache Airflow!

The Airflow community has introduced a limit of 5 open pull requests at a time for contributors without write access to the repository. You currently have 6 open pull requests, so - as a one-time step of introducing the limit - we closed the ones where maintainers have not engaged yet:

These pull requests stay open because maintainers are already engaged in them - they count towards your limit:

This is not a judgement of you or of your changes. We never told contributors before that opening many pull requests at once was a problem, so there is nothing to feel bad about - and nothing is lost: your branches, commits and the review history stay where they are.

What we ask you to do is to make your first prioritization decision: choose which of the pull requests above matter most to you, and reopen them (up to 5 open at a time, including the ones still open) with the "Reopen pull request" button or gh pr reopen <PR_NUMBER> --repo apache/airflow. Reopen the ones you are ready to follow through - keep them rebased, respond to review comments and fix failing checks.

While your pull requests are waiting for review, the most valuable thing you can do is help in other ways - reviewing other contributors' pull requests, helping with issues, and taking part in the discussions on the devlist and Slack.

Why we introduced the limit, what it means for you and how to reopen or restore a pull request is explained in https://github.com/apache/airflow/blob/main/contributing-docs/32_open_pull_request_limit.rst.


Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

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

Labels

area:providers closed because of open PR limit Closed as a one-time step of introducing the open pull request limit kind:documentation provider:qdrant ready for maintainer review Set after triaging when all criteria pass.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants