Skip to content

Instantly share code, notes, and snippets.

@fozziethebeat
Created December 15, 2021 07:35
Show Gist options
  • Save fozziethebeat/4c788b5d329c05d7c427aefc896394a5 to your computer and use it in GitHub Desktop.
Save fozziethebeat/4c788b5d329c05d7c427aefc896394a5 to your computer and use it in GitHub Desktop.
Tries reading some Dataframes and only re-computes new output Dataframes if needed, otherwise copies old result dataframes
import numpy as np
import pandas as pd
import prefect
from os import listdir
from os.path import isfile, join
from prefect import Flow, apply_map, case, task
from prefect.tasks.control_flow import merge
INPUT_BASE_PATH = './data/input'
OUTPUT_BASE_PATH = './data/output'
@task
def emit_filenames():
'''
Returns a list of CSV filen names found in the input data directory.
'''
csvs = [
csv_file for csv_file in listdir('./data/input')
if isfile(join(INPUT_BASE_PATH, csv_file))
]
return csvs
@task
def verify_checkpoint(filename):
'''
Reads the input dataframe and corresponding output dataframe, then figures
out if the checkpoints match.
This assumes the input csv has the following fields:
- region
- metric
- value
- checkpoint
This assumes the output csv has the following fields:
- region
- cases
- deaths
- ratio
- checkpoint
Returns a type of two values:
- True if the checkpoints match, false otherwise
- The output dataframe when checkpoints match, otherwise the input
dataframe
'''
input_path = join(INPUT_BASE_PATH, filename)
output_path = join(OUTPUT_BASE_PATH, filename)
input_dataframe = pd.read_csv(input_path)
output_dataframe = pd.read_csv(output_path)
input_checkpoint = input_dataframe['checkpoint'].min()
output_checkpoint = output_dataframe['checkpoint'].min()
if (input_checkpoint == output_checkpoint):
return (True, output_dataframe)
return (False, input_dataframe)
@task
def checkpoint_matches(dataframes):
'''
Returns the truth value in the tuple
'''
return dataframes[0]
@task
def copy_result(dataframes):
'''
Returns the dataframe in the tuple
'''
return dataframes[1]
@task
def compute_result(dataframe_tuple):
'''
Does some fake computation and copies the input checkpoint to output.
'''
dataframe = dataframe_tuple[1]
checkpoint = dataframe['checkpoint'].min()
result = (dataframe.drop('checkpoint',
axis=1).set_index(['region', 'metric'],
drop=True).unstack())
result.columns = result.columns.droplevel()
result['checkpoint'] = checkpoint
return result
@task
def print_dataframe(dataframe):
'''
Prints the result
'''
print(dataframe)
def compute_or_copy(dataframes):
'''
Copies the old output dataframe if checkpoints match or computes a new
output dataframe if needed.
'''
cond = checkpoint_matches(dataframes)
with case(cond, True):
copy_df = copy_result(dataframes)
with case(cond, False):
compute_df = compute_result(dataframes)
return merge(copy_df, compute_df)
with Flow('ETL') as flow:
csv_filenames = emit_filenames()
dataframes = verify_checkpoint.map(csv_filenames)
final_dataframes = apply_map(compute_or_copy, dataframes)
print_dataframe.map(final_dataframes)
flow.run()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment