cancel future and wait for it

iterative · Jan 1, 2025 · 7adfc0a · 7adfc0a
1 parent 92d0cc5
commit 7adfc0a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/datachain/asyn.py b/src/datachain/asyn.py
@@ -8,7 +8,7 @@
     Iterable,
     Iterator,
 )
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, wait
 from heapq import heappop, heappush
 from typing import Any, Callable, Generic, Optional, TypeVar
 
@@ -179,6 +179,7 @@ def iterate(self, timeout=None) -> Generator[ResultT, None, None]:
             self.shutdown_producer()
             if not async_run.done():
                 async_run.cancel()
+                wait([async_run])
 
     def __iter__(self):
         return self.iterate()