This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit af1f81043b2b9897e0b2f885fc793cecde82bc61 Author: chaoyang <[email protected]> AuthorDate: Sat Sep 12 21:40:50 2026 +0800 [python] Close vector readers when search raises synchronously (#9755) --- .../pypaimon/table/source/vector_search_read.py | 12 ++++- .../pypaimon/tests/vector_search_filter_test.py | 55 ++++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 46c45f01f3..fe9d965367 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -255,7 +255,11 @@ class AbstractVectorSearchReadImpl: reader, offset_reader = self._open_offset_reader( vector_index_files, row_range_start, row_range_end) - future = offset_reader.visit_vector_search(vector_search) + try: + future = offset_reader.visit_vector_search(vector_search) + except BaseException: + reader.close() + raise future.add_done_callback(lambda _: reader.close()) return future @@ -397,7 +401,11 @@ class AbstractVectorSearchReadImpl: reader, offset_reader = self._open_offset_reader( vector_index_files, row_range_start, row_range_end) - future = offset_reader.visit_batch_vector_search(batch_vector_search) + try: + future = offset_reader.visit_batch_vector_search(batch_vector_search) + except BaseException: + reader.close() + raise future.add_done_callback(lambda _: reader.close()) return future diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index e3bbccc911..4d0c39e29e 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -3566,5 +3566,60 @@ class BatchVectorSearchTest(unittest.TestCase): mock.patch.stopall() +class VectorSearchReaderCleanupTest(unittest.TestCase): + + def test_reader_lifetime_for_single_and_batch_search(self): + from concurrent.futures import Future + + from pypaimon.globalindex.offset_global_index_reader import OffsetGlobalIndexReader + from pypaimon.table.source.vector_search_read import DataEvolutionVectorRead + + column = _field(1, "embedding", "FLOAT") + read = DataEvolutionVectorRead(_StubTable([column], []), 1, column, [1.0]) + for batch in (False, True): + for outcome in ("sync_error", "pending_success", "pending_error", "completed"): + with self.subTest(batch=batch, outcome=outcome): + stream = io.BytesIO(b"index") + reader = mock.Mock() + reader.close.side_effect = stream.close + visit = (reader.visit_batch_vector_search if batch + else reader.visit_vector_search) + error = ValueError("Query vector dimension mismatch") + source = Future() + result = [None, None] if batch else None + if outcome == "sync_error": + visit.side_effect = error + else: + visit.return_value = source + if outcome == "completed": + source.set_result(result) + offset = OffsetGlobalIndexReader(reader, 0, 10) + evaluate = read._eval_batch if batch else read._eval + query = [[1.0], [2.0]] if batch else [1.0] + with mock.patch.object(read, "_open_offset_reader", + return_value=(reader, offset)): + if outcome == "sync_error": + with self.assertRaises(ValueError) as raised: + evaluate(0, 10, [object()], query, 1, None) + self.assertIs(error, raised.exception) + else: + future = evaluate(0, 10, [object()], query, 1, None) + if outcome.startswith("pending"): + reader.close.assert_not_called() + self.assertFalse(stream.closed) + if outcome == "pending_error": + source.set_exception(error) + else: + source.set_result(result) + if outcome == "pending_error": + with self.assertRaises(ValueError) as raised: + future.result() + self.assertIs(error, raised.exception) + else: + self.assertEqual(result, future.result()) + reader.close.assert_called_once_with() + self.assertTrue(stream.closed) + + if __name__ == "__main__": unittest.main()
