ari/minerva
1
0
Fork 0
forked from bellwether/minerva
Code Pull requests Activity

added parallel file downloading and dried up the code a touch

Browse source
This commit is contained in:
Ari Brown 2023-08-02 15:12:06 -04:00
parent 68bc346e24
commit 22746b6639
5 changed files with 16 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

6
minerva/__init__.py
View file

@ -1,8 +1,8 @@
from .minerva import Athena, Execute, Query, Redshift, parallel_map
from .parallel import parallel_map
from .athena import Athena
from .redshift import Redshift
__all__ = [
"Execute",
"Query",
"Athena",
"Redshift",
"parallel_map"

8
minerva/athena.py
View file

@ -49,12 +49,12 @@ class Execute:
return self.sql
def run(self):
resp = self.athena.start_query_execution(QueryString=self.sql)
resp = self.athena.start_query_execution(QueryString=self.query())
self.query_id = resp['QueryExecutionId']
return resp
def status(self):
return self.info()['Status']
return self.info()['Status']['State']
def info(self):
res = self.athena.get_query_execution(QueryExecutionId=self.query_id)
@ -62,8 +62,10 @@ class Execute:
return self.info_cache
def finish(self):
while stat := self.status() in ['QUEUED', 'RUNNING']:
stat = self.status()
while stat in ['QUEUED', 'RUNNING']:
time.sleep(5)
stat = self.status()
return stat # finalized state

3
minerva/parallel.py
View file

@ -10,7 +10,8 @@ from joblib import Parallel, delayed
#
# parallel_map(say, [str(i) for i in range(10)], cores=4)
def parallel_map(func=None, data=None, cores=8):
size = len(data) // cores
cores = min(cores, len(data))
size = len(data) // cores
groups = [data[i:i + size] for i in range(0, len(data), size)]
def wrapper_func(fs):

9
minerva/redshift.py
View file

@ -6,6 +6,7 @@ import pyarrow as pa
import pyarrow.dataset
import pprint
import datetime
import json
from minerva import parallel_map
pp = pprint.PrettyPrinter(indent=4)
@ -66,8 +67,10 @@ class Execute:
return self.info_cache
def finish(self):
while stat := self.status() in ['SUBMITTED', 'PICKED', 'STARTED']:
stat = self.status()
while stat in ['SUBMITTED', 'PICKED', 'STARTED']:
time.sleep(5)
stat = self.status()
return stat # finalized state
@ -87,9 +90,9 @@ class Query(Execute):
def manifest_files(self):
status = self.finish()
if status == "SUCCEEDED":
if status == "FINISHED":
# Track the runtime
self.runtime = tiedot['UpdatedAt'] - tiedot['CreatedAt']
self.runtime = self.info_cache['UpdatedAt'] - self.info_cache['CreatedAt']
# Because we're using `UNLOAD`, we get a manifest of the files
# that make up our data.