def test_run_and_write_tables(df, store_name):
orca.add_table('table', df)
def year_key(y):
return '{}'.format(y)
def series_year(y):
return pd.Series([y] * 3, index=df.index, name=str(y))
@orca.step()
def step(iter_var, table):
table[year_key(iter_var)] = series_year(iter_var)
orca.run(
['step'], iter_vars=range(11), data_out=store_name, out_interval=3)
with pd.HDFStore(store_name, mode='r') as store:
for year in range(0, 11, 3):
key = '{}/table'.format(year)
assert key in store
for x in range(year):
pdt.assert_series_equal(
store[key][year_key(x)], series_year(x))
assert 'base/table' in store
for x in range(11):
pdt.assert_series_equal(
store['10/table'][year_key(x)], series_year(x))
def test_step_run(df):
orca.add_table('test_table', df)
@orca.table()
def table_func(test_table):
tt = test_table.to_frame()
tt['c'] = [7, 8, 9]
return tt
@orca.column('table_func')
def new_col(test_table, table_func):
tt = test_table.to_frame()
tf = table_func.to_frame(columns=['c'])
return tt['a'] + tt['b'] + tf['c']
@orca.step()
def test_step1(iter_var, test_table, table_func):
tf = table_func.to_frame(columns=['new_col'])
test_table[iter_var] = tf['new_col'] + iter_var
@orca.step('test_step2')
def asdf(table='test_table'):
tt = table.to_frame()
table['a'] = tt['a']**2
orca.run(steps=['test_step1', 'test_step2'], iter_vars=[2000, 3000])
test_table = orca.get_table('test_table')
assert_frames_equal(
test_table.to_frame(),
pd.DataFrame(
{
'a': [1, 16, 81],
'b': [4, 5, 6],
2000: [2012, 2015, 2018],
3000: [3012, 3017, 3024]
},
index=['x', 'y', 'z']))
m = orca.get_step('test_step1')
assert set(m._tables_used()) == {'test_table', 'table_func'}
def test_injectables_combined(df):
@orca.injectable()
def column():
return pd.Series(['a', 'b', 'c'], index=df.index)
@orca.table()
def table():
return df
@orca.step()
def step(table, column):
df = table.to_frame()
df['new'] = column
orca.add_table('table', df)
orca.run(steps=['step'])
table_wr = orca.get_table('table').to_frame()
pdt.assert_frame_equal(table_wr[['a', 'b']], df)
pdt.assert_series_equal(table_wr['new'], pd.Series(column(), name='new'))
def test_cache_scope(df):
orca.add_injectable('x', 11)
orca.add_injectable('y', 22)
orca.add_injectable('z', 33)
orca.add_injectable('iterations', 1)
@orca.injectable(cache=True, cache_scope='forever')
def a(x):
return x
@orca.injectable(cache=True, cache_scope='iteration')
def b(y):
return y
@orca.injectable(cache=True, cache_scope='step')
def c(z):
return z
@orca.step()
def m1(iter_var, a, b, c):
orca.add_injectable('x', iter_var + a)
orca.add_injectable('y', iter_var + b)
orca.add_injectable('z', iter_var + c)
assert a == 11
@orca.step()
def m2(iter_var, a, b, c, iterations):
assert a == 11
if iter_var == 1000:
assert b == 22
assert c == 1033
elif iter_var == 2000:
assert b == 1022
assert c == 3033
orca.add_injectable('iterations', iterations + 1)
orca.run(['m1', 'm2'], iter_vars=[1000, 2000])
def test_step_run(df):
orca.add_table('test_table', df)
@orca.table()
def table_func(test_table):
tt = test_table.to_frame()
tt['c'] = [7, 8, 9]
return tt
@orca.column('table_func')
def new_col(test_table, table_func):
tt = test_table.to_frame()
tf = table_func.to_frame(columns=['c'])
return tt['a'] + tt['b'] + tf['c']
@orca.step()
def test_step1(iter_var, test_table, table_func):
tf = table_func.to_frame(columns=['new_col'])
test_table[iter_var] = tf['new_col'] + iter_var
@orca.step('test_step2')
def asdf(table='test_table'):
tt = table.to_frame()
table['a'] = tt['a'] ** 2
orca.run(steps=['test_step1', 'test_step2'], iter_vars=[2000, 3000])
test_table = orca.get_table('test_table')
assert_frames_equal(
test_table.to_frame(),
pd.DataFrame(
{'a': [1, 16, 81],
'b': [4, 5, 6],
2000: [2012, 2015, 2018],
3000: [3012, 3017, 3024]},
index=['x', 'y', 'z']))
m = orca.get_step('test_step1')
assert set(m._tables_used()) == {'test_table', 'table_func'}
def test_run_and_write_tables_out_tables_provided(df, store_name):
table_names = ['table', 'table2', 'table3']
for t in table_names:
orca.add_table(t, df)
@orca.step()
def step(iter_var, table, table2):
return
orca.run(['step'],
iter_vars=range(1),
data_out=store_name,
out_base_tables=table_names,
out_run_tables=['table'])
with pd.HDFStore(store_name, mode='r') as store:
for t in table_names:
assert 'base/{}'.format(t) in store
assert '0/table' in store
assert '0/table2' not in store
assert '0/table3' not in store
def test_run_and_write_tables_out_tables_provided(df, store_name):
table_names = ['table', 'table2', 'table3']
for t in table_names:
orca.add_table(t, df)
@orca.step()
def step(iter_var, table, table2):
return
orca.run(
['step'],
iter_vars=range(1),
data_out=store_name,
out_base_tables=table_names,
out_run_tables=['table'])
with pd.HDFStore(store_name, mode='r') as store:
for t in table_names:
assert 'base/{}'.format(t) in store
assert '0/table' in store
assert '0/table2' not in store
assert '0/table3' not in store