From 575aaca4378bce7b56fd6defb99c9ef7fd909c66 Mon Sep 17 00:00:00 2001 From: harbinresearcher Date: Sun, 11 Oct 2026 14:35:25 +0800 Subject: [PATCH] Release partition_unique metadata references correctly --- streamz/core.py | 5 ++++- streamz/tests/test_core.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/streamz/core.py b/streamz/core.py index 6d49d5a8..b7bc8905 100644 --- a/streamz/core.py +++ b/streamz/core.py @@ -1249,16 +1249,19 @@ def update(self, x, who=None, metadata=None): # remove key if already present so that emitted value # will reflect elements' actual relative ordering self._buffer.pop(y, None) - self._metadata_buffer.pop(y, None) + self._release_refs(self._metadata_buffer.pop(y, [])) self._buffer[y] = x self._metadata_buffer[y] = metadata else: # self.keep == "first" if y not in self._buffer: self._buffer[y] = x self._metadata_buffer[y] = metadata + else: + self._release_refs(metadata) if len(self._buffer) == self.n: result, self._buffer = tuple(self._buffer.values()), {} metadata_result, self._metadata_buffer = list(self._metadata_buffer.values()), {} + metadata_result = [m for ml in metadata_result for m in ml] ret = self._emit(result, metadata_result) self._release_refs(metadata_result) return ret diff --git a/streamz/tests/test_core.py b/streamz/tests/test_core.py index 67261255..4d0cfb1a 100644 --- a/streamz/tests/test_core.py +++ b/streamz/tests/test_core.py @@ -274,6 +274,36 @@ def test_partition_unique(n, key, keep, elements, exp_result): assert L == exp_result +@pytest.mark.parametrize("keep", ["first", "last"]) +def test_partition_unique_ref_counts(keep): + source = Stream() + partition = source.partition_unique(2, keep=keep) + collected = partition.collect() + refs = [RefCounter() for _ in range(3)] + + source.emit(1, metadata=[{'ref': refs[0]}]) + source.emit(1, metadata=[{'ref': refs[1]}]) + expected = [1, 0, 0] if keep == "first" else [0, 1, 0] + assert [ref.count for ref in refs] == expected + + source.emit(2, metadata=[{'ref': refs[2]}]) + expected[2] = 1 + assert [ref.count for ref in refs] == expected + collected.flush() + assert all(ref.count == 0 for ref in refs) + + +@pytest.mark.parametrize("keep", ["first", "last"]) +def test_partition_unique_metadata(keep): + source = Stream() + output = metadata(source.partition_unique(2, keep=keep)).sink_to_list() + source.emit(1, metadata=[{'v': 0}, {'v': 1}]) + source.emit(1, metadata=[{'v': 2}]) + source.emit(2, metadata=[{'v': 3}]) + retained = [{'v': 0}, {'v': 1}] if keep == "first" else [{'v': 2}] + assert output == [retained + [{'v': 3}]] + + def test_partition_timeout(): source = Stream() L = source.partition(10, timeout=0.01).sink_to_list()