This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new de56e6568bf [Python] Optionally write chunk content to Qdrant payload
(#40502)
de56e6568bf is described below
commit de56e6568bfa7cab0daf0dfa49dded1bdda842e6
Author: Wangshu <[email protected]>
AuthorDate: Thu Oct 15 06:56:56 2026 -0700
[Python] Optionally write chunk content to Qdrant payload (#40502)
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 <[email protected]>
---
CHANGES.md | 1 +
sdks/python/apache_beam/ml/rag/ingestion/qdrant.py | 17 +++++-
.../apache_beam/ml/rag/ingestion/qdrant_it_test.py | 25 +++++++++
.../apache_beam/ml/rag/ingestion/qdrant_test.py | 64 ++++++++++++++++++++++
4 files changed, 106 insertions(+), 1 deletion(-)
diff --git a/CHANGES.md b/CHANGES.md
index 14c7e676f18..c787b83dde7 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)).
* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g.
`--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set
defaults for all AWS service clients. The existing `--s3_*` options in
`S3Options` override these defaults for the S3 client only
([#40528](https://github.com/apache/beam/issues/40528)).
* (Python) Added `--aws_credentials_secret_manager` so AWS credentials
(`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and
their `--s3_*` overrides) can be secret specs resolved through a secret manager
such as `GoogleCloudSecretManager`, instead of plain-text values
([#40528](https://github.com/apache/beam/issues/40528)).
diff --git a/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py
b/sdks/python/apache_beam/ml/rag/ingestion/qdrant.py
index abe9efc56cb..a1e3fd7c8cf 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 @@ class QdrantWriteConfig(VectorDatabaseWriteConfig):
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 b82ce2b6e74..7148b222932 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 @@ class TestQdrantIngestion(unittest.TestCase):
)
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 ff4ee14e97a..f969b09cb56 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 @@ class TestQdrantCreateConverter(unittest.TestCase):
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):
