Add QdrantSearchOperator for vector similarity search#69673
Open
YAshhh29 wants to merge 6 commits into
Open
Conversation
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Qdrant provider today lets users write vectors into a collection
(
QdrantIngestOperator) but has no operator for the other half of a RAGpipeline: reading them back out. Users who want to run a similarity search
from a DAG have to reach into
hook.conn.query_pointsdirectly and rememberto convert the returned pydantic
ScoredPointobjects into plain dicts soAirflow 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 thatgap. 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
searchmethod 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:
0 open issues mentioning "qdrant search" -- greenfield, no one else
was building this.
qdrant-client 1.18.0(the provider pins
>=1.17.1) exposesquery_pointswith every oneof the 9 named parameters this operator forwards; the older
search()method is deprecated and slated for removal.
QueryResponse.pointsis
List[ScoredPoint], andScoredPoint.model_dump()produces theid/score/payload/vector/version/shard_key/order_valuedict shapethe operator promises callers.
python-modulesinprovider.yamlcovers any class inoperators/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
QdrantHook.search()wrapsQdrantClient.query_pointsand convertseach returned
ScoredPointto a plaindictviamodel_dump(). Theoperator 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.
list[dict[str, Any]](id, score, payload, and optionally vector), so results land in XCom
without any user-side workaround.
query_points, not the deprecatedsearch(). Thesearch()API in
qdrant-clientis scheduled for removal in a future major;query_pointsis the modern surface (also supports named/sparsevectors, hybrid search, etc.). A regression test asserts we never
fall back to the deprecated method.
**kwargspassthrough. Forwards anyquery_pointsparameter wedon't enumerate (
using,prefetch,lookup_from, ...) so the hookstays forward-compatible with hybrid search and named vectors without
a follow-up PR.
What changes
providers/qdrant/src/airflow/providers/qdrant/hooks/qdrant.pyQdrantHook.search(...)method wrappingquery_points, returninglist[dict]viaScoredPoint.model_dump().providers/qdrant/src/airflow/providers/qdrant/operators/qdrant.pyQdrantSearchOperatorclass alongside the existingQdrantIngestOperator.template_fieldsincludecollection_name,query,query_filter,limitso a RAG DAG can XCom-pull a queryvector from an upstream embedding task.
providers/qdrant/tests/unit/qdrant/hooks/test_qdrant.pylist[dict]viamodel_dump; usesquery_points(not deprecatedsearch) with all named argsforwarded; extra
**kwargsalso forwarded.providers/qdrant/tests/unit/qdrant/operators/test_qdrant.pyexpected; every optional arg reaches the hook; template_fields
cover the runtime parameters; default
conn_idmatches the hook's.providers/qdrant/tests/system/qdrant/example_dag_qdrant.pyQdrantSearchOperatortask downstream of the existingingest task with
# [START/END] howto_operator_qdrant_searchmarkers.
providers/qdrant/docs/operators/qdrant.rst.. _howto/operator:QdrantSearchOperator:anchor and an
.. exampleinclude::pulling the DAG snippet.No
provider.yaml/get_provider_info.pychanges needed: the registrylists 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 logby the release manager perAGENTS.md.Testing
a standalone harness that runs the real hook/operator code with mocked
Qdrant client + a
BaseHook/BaseOperatorshim (full Airflow can'trun on Windows).
query_pointsaccepts every one of the 9 named parameters we forward, and
QueryResponse.pointsis aList[ScoredPoint]with the expectedmodel_dump()shape (id,score,payload,vector, ...).QdrantIngestOperatorstill constructs andbehaves identically -- we only added to the module, no existing code
was touched.
ruff check+ruff format --check: 26 files clean..github/instructions/code-review.instructions.md-- no red flags (no
time.time, noassertin prod, no newAirflowException, no British spellings, no missing tests).Was generative AI tooling used to co-author this PR?
Generated-by: GitHub Copilot (Claude Opus 4.6) following the guidelines