diff --git a/packages/bigframes/bigframes/session/anonymous_dataset.py b/packages/bigframes/bigframes/session/anonymous_dataset.py index 1a3d43655b79..ed718ff909f0 100644 --- a/packages/bigframes/bigframes/session/anonymous_dataset.py +++ b/packages/bigframes/bigframes/session/anonymous_dataset.py @@ -16,6 +16,7 @@ import threading import uuid import warnings +from concurrent.futures import ThreadPoolExecutor from typing import List, Optional, Sequence import google.cloud.bigquery as bigquery @@ -170,9 +171,19 @@ def _cleanup_old_udfs(self): def close(self): """Delete tables that were created with this session's session_id.""" - for table_ref in self._table_ids: - self.bqclient.delete_table(table_ref, not_found_ok=True) - self._table_ids.clear() + if self._table_ids: + try: + with ThreadPoolExecutor() as executor: + futures = [ + executor.submit( + self.bqclient.delete_table, table_ref, not_found_ok=True + ) + for table_ref in self._table_ids + ] + for future in futures: + future.result() + finally: + self._table_ids.clear() try: # Before closing the session, attempt to clean up any uncollected, diff --git a/packages/bigframes/noxfile.py b/packages/bigframes/noxfile.py index e7c105a552e8..0a33264aa8ea 100644 --- a/packages/bigframes/noxfile.py +++ b/packages/bigframes/noxfile.py @@ -354,6 +354,7 @@ def run_system( "py.test", "-v", f"-n={num_workers}", + "--dist=worksteal", # Any individual test taking longer than 15 mins will be terminated. f"--timeout={timeout_seconds}", # Log 20 slowest tests