Skip to content
This repository has been archived by the owner on Jan 7, 2024. It is now read-only.

Commit

Permalink
⚡ Fix when to validate
Browse files Browse the repository at this point in the history
  • Loading branch information
nagomiso committed Nov 16, 2022
1 parent 31c66de commit 9b27cae
Showing 1 changed file with 12 additions and 12 deletions.
24 changes: 12 additions & 12 deletions src/saori/athena.py
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,18 @@ def save_sql_query_results(
DataFrame | Iterator[DataFrame]: \
Pandas DataFrame or Generator of Pandas DataFrames if chunksize is passed.
"""
default_max_workers = min(32, mp.cpu_count() + 4)
if not max_workers:
max_workers_ = default_max_workers
elif isinstance(max_workers, int):
if max_workers < 0:
max_workers_ = max(default_max_workers + max_workers, 1)
else:
max_workers_ = max_workers
else:
raise ValueError(
f"`max_workers` must be integer. {repr(max_workers)} is not integer."
)
if Path(save_dir).protocol == "file":
if not s3_output:
raise ValueError(
Expand Down Expand Up @@ -394,18 +406,6 @@ def save_sql_query_results(
next(query_resuts) # type: ignore
move_file = partial(_move_file, dst=save_dir)
sources = tempdir.glob("temp_table*/*")
default_max_workers = min(32, mp.cpu_count() + 4)
if not max_workers:
max_workers_ = default_max_workers
elif isinstance(max_workers, int):
if max_workers < 0:
max_workers_ = max(default_max_workers + max_workers, 1)
else:
max_workers_ = max_workers
else:
raise ValueError(
f"`max_workers` must be integer. {repr(max_workers)} is not integer."
)
if 1 < max_workers_:
with ThreadPoolExecutor(max_workers=max_workers_) as executor:
for _ in executor.map(move_file, sources):
Expand Down

0 comments on commit 9b27cae

Please sign in to comment.