From 30227d42110b394d69f8bc0193f66a0a428df99e Mon Sep 17 00:00:00 2001 From: Wangshu Pang Date: Sat, 10 Oct 2026 14:19:23 -0700 Subject: [PATCH] [Python] Optionally write chunk content to Qdrant payload Add QdrantWriteConfig.content_payload_key. When set, the converter stores EmbeddableItem.content_string under that key in the point payload next to the metadata, so points retrieved from a collection written by Beam carry the chunk text. The default (None) keeps the current behavior of writing only the metadata. A key that collides with a metadata key raises ValueError instead of overwriting it, an empty key is rejected at construction time, and the item's metadata dict is copied rather than mutated. Fixes #40476 Co-Authored-By: Claude Opus 5.5 --- CHANGES.md | 1 + .../apache_beam/ml/rag/ingestion/qdrant.py | 17 ++++- .../ml/rag/ingestion/qdrant_it_test.py | 25 ++++++++ .../ml/rag/ingestion/qdrant_test.py | 64 +++++++++++++++++++ 4 files changed, 106 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..b4fa58857f4d 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -72,6 +72,7 @@ * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)). * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)). +* `QdrantWriteConfig` can now optionally store chunk content in the point payload via `content_payload_key`, so retrieved points include the chunk text (Python) ([#40476](https://github.com/apache/beam/issues/40476)). * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). ## Breaking Changes diff --git a/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py b/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py index abe9efc56cbf..a1e3fd7c8cf3 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py @@ -183,6 +183,11 @@ class QdrantWriteConfig(VectorDatabaseWriteConfig): kwargs: Additional keyword arguments to pass to the client's upsert method. dense_embedding_key: name for the dense vector in the qdrant collection. sparse_embedding_key: name for the sparse vector in the qdrant collection. + content_payload_key: Optional payload key under which the item's content + (``EmbeddableItem.content_string``) is stored, so that the chunk text is + returned with search results. If ``None`` (default), content is not + written and only ``metadata`` is stored as the payload. A ``ValueError`` + is raised at conversion time if the key collides with a metadata key. """ connection_params: QdrantConnectionParameters @@ -193,12 +198,15 @@ class QdrantWriteConfig(VectorDatabaseWriteConfig): kwargs: dict[str, Any] = field(default_factory=dict) dense_embedding_key: str = "dense" sparse_embedding_key: str = "sparse" + content_payload_key: Optional[str] = None def __post_init__(self): if not self.collection_name: raise ValueError("Collection name must be provided") if self.batch_size <= 0: raise ValueError("Batch size must be a positive integer") + if self.content_payload_key is not None and not self.content_payload_key: + raise ValueError("content_payload_key must be a non-empty string") def create_write_transform(self) -> beam.PTransform[EmbeddableItem, Any]: return _QdrantWriteTransform(self) @@ -222,10 +230,17 @@ def convert(item: EmbeddableItem) -> "models.PointStruct": id = ( int(item.id) if isinstance(item.id, str) and item.id.isdigit() else item.id) + payload = dict(item.metadata) if item.metadata else {} + if self.content_payload_key is not None: + if self.content_payload_key in payload: + raise ValueError( + f"content_payload_key '{self.content_payload_key}' collides " + f"with an existing metadata key for item {item.id}") + payload[self.content_payload_key] = item.content_string return models.PointStruct( id=id, vector=vector, - payload=item.metadata if item.metadata else None, + payload=payload if payload else None, ) return convert diff --git a/sdks/python/apache_beam/ml/rag/ingestion/qdrant_it_test.py b/sdks/python/apache_beam/ml/rag/ingestion/qdrant_it_test.py index b82ce2b6e74d..7148b2229328 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/qdrant_it_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/qdrant_it_test.py @@ -134,6 +134,31 @@ def test_write_dense_embeddings_only(self): ) self.assertEqual(expected_record, points_by_id[int(item.id)]) + def test_write_with_content_payload_key(self): + write_config = QdrantWriteConfig( + connection_params=self._connection_params, + collection_name=self._collection_name, + batch_size=len(TEST_CORPUS), + content_payload_key="content", + ) + + with TestPipeline(is_integration_test=True) as p: + _ = p | beam.Create(TEST_CORPUS) | write_config.create_write_transform() + + points, _ = self.client.scroll( + collection_name=self._collection_name, + limit=100, + with_payload=True, + with_vectors=False, + ) + points_by_id = {p.id: p for p in points} + self.assertEqual(len(points_by_id), len(TEST_CORPUS)) + + for item in TEST_CORPUS: + expected_payload = dict(item.metadata or {}) + expected_payload["content"] = item.content_string + self.assertEqual(expected_payload, points_by_id[int(item.id)].payload) + def test_write_sparse_embeddings_only(self): sparse_corpus = [ EmbeddableItem( diff --git a/sdks/python/apache_beam/ml/rag/ingestion/qdrant_test.py b/sdks/python/apache_beam/ml/rag/ingestion/qdrant_test.py index ff4ee14e97a0..f969b09cb56b 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/qdrant_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/qdrant_test.py @@ -269,6 +269,70 @@ def test_convert_from_text_factory(self): self.assertIsInstance(result, models.PointStruct) self.assertIn("dense", result.vector) + def test_content_not_in_payload_by_default(self): + item = EmbeddableItem( + id="1", + content=Content(text="hello"), + embedding=Embedding(dense_embedding=[1.0]), + metadata={"source": "test"}, + ) + result = self.convert(item) + self.assertEqual(result.payload, {"source": "test"}) + + def test_content_payload_key_writes_content(self): + config = QdrantWriteConfig( + connection_params=QdrantConnectionParameters(location=":memory:"), + collection_name="test", + content_payload_key="content", + ) + item = EmbeddableItem( + id="1", + content=Content(text="hello"), + embedding=Embedding(dense_embedding=[1.0]), + metadata={"source": "test"}, + ) + result = config.create_converter()(item) + self.assertEqual(result.payload, {"source": "test", "content": "hello"}) + # metadata of the input item must not be mutated + self.assertEqual(item.metadata, {"source": "test"}) + + def test_content_payload_key_without_metadata(self): + config = QdrantWriteConfig( + connection_params=QdrantConnectionParameters(location=":memory:"), + collection_name="test", + content_payload_key="text", + ) + item = EmbeddableItem( + id="1", + content=Content(text="hello"), + embedding=Embedding(dense_embedding=[1.0]), + ) + result = config.create_converter()(item) + self.assertEqual(result.payload, {"text": "hello"}) + + def test_content_payload_key_collision_raises(self): + config = QdrantWriteConfig( + connection_params=QdrantConnectionParameters(location=":memory:"), + collection_name="test", + content_payload_key="source", + ) + item = EmbeddableItem( + id="1", + content=Content(text="hello"), + embedding=Embedding(dense_embedding=[1.0]), + metadata={"source": "test"}, + ) + with self.assertRaises(ValueError): + config.create_converter()(item) + + def test_empty_content_payload_key_raises(self): + with self.assertRaises(ValueError): + QdrantWriteConfig( + connection_params=QdrantConnectionParameters(location=":memory:"), + collection_name="test", + content_payload_key="", + ) + @unittest.skipIf(not QDRANT_AVAILABLE, "qdrant dependencies not installed.") class TestQdrantWriteFnBatching(unittest.TestCase):