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 bef9ef19056 [Python] Reuse Milvus client in RAG vector sink write path
(#39378)
bef9ef19056 is described below
commit bef9ef19056632ed3a15e60ba5202478a668ac80
Author: Venkata Bharath Malapati <[email protected]>
AuthorDate: Sun Jul 19 15:35:12 2026 -0400
[Python] Reuse Milvus client in RAG vector sink write path (#39378)
* Reuse Milvus client in vector sink write path
_MilvusSink.write() constructed a new MilvusClient on every call, which
discarded the retry-wrapped client created in __enter__. This bypassed the
connection retry/backoff and left the initial client unclosed (leaked),
while __exit__ only closed the last client. The write() docstring also
incorrectly claimed the collection was flushed for durability.
Reuse the client established in __enter__ for the upsert, and correct the
docstring to describe the actual (no-flush) visibility semantics. Add a
unit test asserting write() reuses the __enter__ client instead of
reconstructing one per call.
Co-authored-by: Cursor <[email protected]>
* Clear Milvus client after close in sink __exit__
Set self._client to None after closing so a reused sink instance does
not keep a stale closed client. Address review feedback.
Co-authored-by: Cursor <[email protected]>
---------
Co-authored-by: Cursor <[email protected]>
---
.../apache_beam/ml/rag/ingestion/milvus_search.py | 13 ++++----
.../ml/rag/ingestion/milvus_search_test.py | 35 ++++++++++++++++++++++
2 files changed, 43 insertions(+), 5 deletions(-)
diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py
b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py
index b7cad3796b1..cfca29a1ada 100644
--- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py
+++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search.py
@@ -287,16 +287,18 @@ class _MilvusSink:
"""Writes a batch of documents to the Milvus collection.
Performs an upsert operation to insert new documents or update existing
- ones based on primary key. After the upsert, flushes the collection to
- ensure data persistence.
+ ones based on primary key, reusing the client established in ``__enter__``.
+
+ Note:
+ This submits the upsert to Milvus; it does not call ``flush()`` on the
+ collection. Newly written rows become queryable according to Milvus'
+ normal segment-sync behavior. Callers that require immediate
+ read-after-write visibility should flush the collection explicitly.
Args:
documents: List of dictionaries representing Milvus records to write.
Each dictionary should contain fields matching the collection schema.
"""
- self._client = MilvusClient(
- **unpack_dataclass_with_kwargs(self._connection_params))
-
resp = self._client.upsert(
collection_name=self._write_config.collection_name,
partition_name=self._write_config.partition_name,
@@ -346,3 +348,4 @@ class _MilvusSink:
_ = exc_type, exc_val, exc_tb # Unused parameters
if self._client:
self._client.close()
+ self._client = None
diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py
b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py
index 80d55ac9382..0c67e15823e 100644
--- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py
+++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_test.py
@@ -15,12 +15,15 @@
# limitations under the License.
#
import unittest
+from unittest import mock
from parameterized import parameterized
try:
+ from apache_beam.ml.rag.ingestion import milvus_search as
milvus_search_module
from apache_beam.ml.rag.ingestion.milvus_search import
MilvusVectorWriterConfig
from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig
+ from apache_beam.ml.rag.ingestion.milvus_search import _MilvusSink
from apache_beam.ml.rag.utils import MilvusConnectionParameters
except ImportError as e:
raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}')
@@ -119,5 +122,37 @@ class TestMilvusVectorWriterConfig(unittest.TestCase):
self.assertIn(expected_error_msg, str(context.exception))
+class TestMilvusSinkClientReuse(unittest.TestCase):
+ """Unit tests for Milvus sink client reuse.
+
+ Verifies that ``_MilvusSink.write`` reuses the client established in
+ ``__enter__`` instead of constructing (and leaking) a new client on every
+ write.
+ """
+ def _sink(self):
+ connection_params =
MilvusConnectionParameters(uri="http://localhost:19530";)
+ write_config = MilvusWriteConfig(collection_name="test_collection")
+ return _MilvusSink(connection_params, write_config)
+
+ def test_write_reuses_enter_client(self):
+ """write() must not construct a new MilvusClient per call."""
+ with mock.patch.object(milvus_search_module,
+ 'MilvusClient') as mock_client_cls:
+ with self._sink() as sink:
+ # One client created on __enter__.
+ self.assertEqual(mock_client_cls.call_count, 1)
+
+ sink.write([{"id": 1}])
+ sink.write([{"id": 2}])
+
+ # write() reuses that client; no new construction.
+ self.assertEqual(mock_client_cls.call_count, 1)
+ self.assertEqual(mock_client_cls.return_value.upsert.call_count, 2)
+
+ # Client closed once on __exit__, and cleared so the sink can be reused.
+ self.assertEqual(mock_client_cls.return_value.close.call_count, 1)
+ self.assertIsNone(sink._client)
+
+
if __name__ == '__main__':
unittest.main()
