From 637c1128504014fe32e49f9997c607ef48456d9e Mon Sep 17 00:00:00 2001
From: Mola-maker <229602071+Mola-maker@users.noreply.github.com>
Date: Mon, 12 Oct 2026 20:51:52 +0800
Subject: [PATCH] Fix DaskRunner dropping finish_bundle outputs (#40519)

---
 CHANGES.md                                    |  1 +
 .../runners/dask/dask_runner_test.py          | 35 +++++++++++++++++++
 .../runners/dask/transform_evaluator.py       |  3 +-
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 6b01d163823c..cf47dbfde715 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -89,6 +89,7 @@
 
 ## Bugfixes
 
+* (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)).
 * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)).
 * (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
 * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
diff --git a/sdks/python/apache_beam/runners/dask/dask_runner_test.py b/sdks/python/apache_beam/runners/dask/dask_runner_test.py
index e1e5a4403b46..160a8f3cc51a 100644
--- a/sdks/python/apache_beam/runners/dask/dask_runner_test.py
+++ b/sdks/python/apache_beam/runners/dask/dask_runner_test.py
@@ -125,6 +125,41 @@ def double(x):
       pcoll = p | beam.Create(list(range(1, 11))) | beam.Map(double)
       assert_that(pcoll, equal_to(list(range(2, 21, 2))))
 
+  def test_pardo_finish_bundle_outputs(self):
+    class BufferElements(beam.DoFn):
+      def start_bundle(self):
+        self.elements = []
+
+      def process(self, element):
+        self.elements.append(element)
+        return []
+
+      def finish_bundle(self):
+        for element in self.elements:
+          yield window.WindowedValue(element, 0, [window.GlobalWindow()])
+
+    with self.pipeline as p:
+      pcoll = p | beam.Create([1, 2, 3]) | beam.ParDo(BufferElements())
+      assert_that(pcoll, equal_to([1, 2, 3]))
+
+  def test_pardo_process_and_finish_bundle_outputs(self):
+    class EmitAtProcessAndFinish(beam.DoFn):
+      def start_bundle(self):
+        self.elements = []
+
+      def process(self, element):
+        self.elements.append(element)
+        yield element
+
+      def finish_bundle(self):
+        for element in self.elements:
+          yield window.WindowedValue(-element, 0, [window.GlobalWindow()])
+
+    with self.pipeline as p:
+      pcoll = (
+          p | beam.Create([1, 2, 3]) | beam.ParDo(EmitAtProcessAndFinish()))
+      assert_that(pcoll, equal_to([1, 2, 3, -1, -2, -3]))
+
   def test_create_map_and_groupby(self):
     def double(x):
       return x * 2, x
diff --git a/sdks/python/apache_beam/runners/dask/transform_evaluator.py b/sdks/python/apache_beam/runners/dask/transform_evaluator.py
index dbf55a3cee0d..43ac488b2937 100644
--- a/sdks/python/apache_beam/runners/dask/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/dask/transform_evaluator.py
@@ -195,9 +195,10 @@ def apply_dofn_to_bundle(
   for it in items:
     do_fn_invoker.invoke_process(it)
 
+  do_fn_invoker.invoke_finish_bundle()
+
   results = [v.value for v in tagged_receivers.values]
 
-  do_fn_invoker.invoke_finish_bundle()
   do_fn_invoker.invoke_teardown()
 
   return results