def test_put_dataframe(self):
# create some dataframe
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
df2 = store.get('mydata')
self.assertTrue(df.equals(df2), "expected dataframes to be equal")
def test_store_with_metadata(self):
om = OmegaStore(prefix='')
# dict
data = {
'a': list(range(1, 10)),
'b': list(range(1, 10))
}
attributes = {'foo': 'bar'}
meta = om.put(data, 'data', attributes=attributes)
self.assertEqual(meta.kind, 'python.data')
self.assertEqual(meta.attributes, attributes)
data2 = om.get('data')
self.assertEqual([data], data2)
# dataframe
df = pd.DataFrame(data)
meta = om.put(df, 'datadf', attributes=attributes)
self.assertEqual(meta.kind, 'pandas.dfrows')
self.assertEqual(meta.attributes, attributes)
df2 = om.get('datadf')
assert_frame_equal(df, df2)
# model
lr = LogisticRegression(solver='liblinear', multi_class='auto')
meta = om.put(lr, 'mymodel', attributes=attributes)
self.assertEqual(meta.kind, 'sklearn.joblib')
self.assertEqual(meta.attributes, attributes)
lr2 = om.get('mymodel')
self.assertIsInstance(lr2, LogisticRegression)
def test_store_irregular_column_names(self):
""" test storing irregular column names """
df = pd.DataFrame({'x_1': range(10)})
store = OmegaStore()
store.put(df, 'foo', append=False)
df2 = store.get('foo')
self.assertEqual(df.columns, df2.columns)
def test_migrate_unhashed_name(self):
store = OmegaStore(bucket='foo', prefix='foo/')
df = pd.DataFrame({'x': range(100)})
long_name = 'a' * 10
raised = False
error = ''
# save as unhashed (old version)
store.defaults.OMEGA_STORE_HASHEDNAMES = False
meta_unhashed = store.put(df, long_name)
# simulate upgrade, no migration
store.defaults.OMEGA_STORE_HASHEDNAMES = True
# check we can still retrieve
dfx = store.get(long_name)
assert_frame_equal(df, dfx)
# migrate
store.defaults.OMEGA_STORE_HASHEDNAMES = True
migrate_unhashed_datasets(store)
meta_migrated = store.metadata(long_name)
# check we can still retrieve after migration
dfx = store.get(long_name)
assert_frame_equal(df, dfx)
# stored hashed
meta_hashed = store.put(df, long_name, append=False)
# check migration worked as expected
self.assertNotEqual(meta_unhashed.collection, meta_hashed.collection)
self.assertEqual(meta_migrated.collection, meta_hashed.collection)
def test_get_dataframe_colspec_opspec(self):
# create some dataframe
df = pd.DataFrame({
'a': list(range(1, 10)),
'b': list(range(1, 10)),
'c': list(range(1, 10)),
})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
# check we can specify [] and # qualifiers
value = store.get('mydata[a]#')
self.assertTrue(hasattr(value, '__next__'))
nvalue = next(value)
self.assertEqual(len(nvalue), len(df))
assert_frame_equal(nvalue, df[['a']])
# check we can specify specific operator
value = store.get('mydata[a,b]#iterchunks')
nvalue = next(value)
self.assertTrue(hasattr(value, '__next__'))
self.assertEqual(len(nvalue), len(df))
assert_frame_equal(nvalue, df[['a', 'b']])
# check we can specify kwargs
value = store.get('mydata[a,b]#iterchunks:chunksize=1')
nvalue = next(value)
self.assertTrue(hasattr(value, '__next__'))
self.assertEqual(len(nvalue), 1)
assert_frame_equal(nvalue, df[['a', 'b']].iloc[0:1])
def test_get_dataframe_projected_mixin(self):
# create some dataframe
df = pd.DataFrame({
'a': list(range(1, 10)),
'b': list(range(1, 10)),
'c': list(range(1, 10)),
})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
# filter in mongodb
specs = ['a', ':b', ':', 'b:', '^c']
for spec in specs:
name_spec = 'mydata[{}]'.format(spec)
df2 = store.get(name_spec)
# filter local dataframe
if spec == ':':
dfx = df.loc[:, :]
elif ':' in spec:
from_col, to_col = spec.split(':')
slice_ = slice(from_col or None, to_col or None)
dfx = df.loc[:, slice_]
elif spec.startswith('^'):
spec_cols = spec[1:].split(',')
cols = [col for col in df.columns if col not in spec_cols]
dfx = df[cols]
else:
dfx = df[[spec]]
self.assertTrue(dfx.equals(df2), "expected dataframes to be equal")
def test_put_python_dict(self):
# create some data
data = {'a': list(range(1, 10)), 'b': list(range(1, 10))}
store = OmegaStore(prefix='')
store.put(data, 'mydata')
data2 = store.get('mydata')
self.assertEquals([data], data2)
def test_put_dataframe_with_index(self):
# create some dataframe
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
store.put(df, 'mydata', index=['a', '-b'])
idxs = list(store.collection('mydata').list_indexes())
idx_names = map(lambda v: dict(v).get('name'), idxs)
self.assertIn('asc_a__desc_b', idx_names)
def test_store_series(self):
""" test storing a pandas series with it's own index """
from string import ascii_lowercase
series = pd.Series(range(10), index=(c for c in ascii_lowercase[0:10]))
store = OmegaStore()
store.put(series, 'fooseries', append=False)
series2 = store.get('fooseries')
assert_series_equal(series, series2)
def test_store_datetime(self):
""" test storing naive datetimes """
df = pd.DataFrame(
{'x': pd.date_range(datetime(2016, 1, 1), datetime(2016, 1, 10))})
store = OmegaStore()
store.put(df, 'test-date', append=False)
df2 = store.get('test-date')
assert_frame_equal(df, df2)
def test_store_dict_in_df(self):
df = pd.DataFrame({
'x': [{'foo': 'bar '}],
})
store = OmegaStore()
store.put(df, 'test-dict', append=False)
df2 = store.get('test-dict')
testing.assert_frame_equal(df, df2)
def test_hidden_temp_handling(self):
foo_store = OmegaStore(bucket='foo')
foo_store.put({}, '_temp')
self.assertNotIn('_temp', foo_store.list(include_temp=False))
self.assertIn('_temp', foo_store.list(include_temp=True))
foo_store.put({}, '.hidden')
self.assertNotIn('.hidden', foo_store.list(hidden=False))
self.assertIn('.hidden', foo_store.list(hidden=True))
def test_store_tz_datetime(self):
""" test storing timezoned datetimes """
df = pd.DataFrame({
'y': pd.date_range('2019-10-01', periods=5, tz='US/Eastern', normalize=True)
})
store = OmegaStore()
store.put(df, 'test-date', append=False)
df2 = store.get('test-date')
testing.assert_frame_equal(df, df2)
def test_put_dataframe_with_index(self):
# create some dataframe
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
store.put(df, 'mydata', index=['a', '-b'])
idxs = store.collection('mydata').index_information()
idx_names = humanize_index(idxs)
self.assertIn('asc__id_asc_a_desc_b_asc__idx#0_0_asc__om#rowid',
idx_names)
def test_get_dataframe_filter(self):
# create some dataframe
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
# filter in mongodb
df2 = store.get('mydata', filter=dict(a__gt=1, a__lt=10))
# filter local dataframe
df = df[(df.a > 1) & (df.a < 10)]
self.assertTrue(df.equals(df2), "expected dataframes to be equal")
def test_get_dataframe_project(self):
# create some dataframe
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
# filter in mongodb
df2 = store.get('mydata', columns=['a'])
# filter local dataframe
df = df[['a']]
self.assertTrue(df.equals(df2), "expected dataframes to be equal")
def test_prefix_store(self):
"""
this is to test if store prefixes work
"""
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
datasets = OmegaStore(prefix='teststore')
models = OmegaStore(prefix='models', kind=Metadata.SKLEARN_JOBLIB)
datasets.put(df, 'test')
self.assertEqual(len(datasets.list()), 1)
self.assertEqual(len(models.list()), 0)
def test_put_python_dict_multiple(self):
# create some data
data = {'a': list(range(1, 10)), 'b': list(range(1, 10))}
store = OmegaStore(prefix='')
store.put(data, 'mydata')
store.put(data, 'mydata')
data2 = store.get('mydata')
# we will have stored the same object twice
self.assertEquals(data, data2[0])
self.assertEquals(data, data2[1])
def test_store_series_timeindex(self):
""" test storing a pandas series with it's own index """
series = pd.Series(range(10),
name='foo',
index=pd.date_range(pd.datetime(2016, 1, 1),
pd.datetime(2016, 1, 10)))
store = OmegaStore()
store.put(series, 'fooseries', append=False)
series2 = store.get('fooseries')
assert_series_equal(series, series2)
def test_put_dataframe_xtra_large(self):
# create some dataframe
# force fast insert
df = pd.DataFrame({
'a': list(range(0, int(1e4 + 1))),
'b': list(range(0, int(1e4 + 1)))
})
store = OmegaStore(prefix='')
store.put(df, 'mydata')
df2 = store.get('mydata')
self.assertTrue(df.equals(df2), "expected dataframes to be equal")
def test_help_docs(self):
foo_store = OmegaStore(bucket='foo')
reg = LinearRegression()
foo_store.put(reg,
'regmodel',
attributes={'docs': 'this is some text'})
# get backend for different signatures
backend = foo_store._resolve_help_backend('regmodel')
self.assertEqual(backend.__doc__, 'this is some text')
backend = foo_store._resolve_help_backend('regmodel', raw=True)
self.assertIsInstance(backend, ScikitLearnBackend)
def test_bucket(self):
# test different buckets actually separate objects by the same name
# -- data
foo_store = OmegaStore(bucket='foo')
bar_store = OmegaStore(bucket='bar')
foo_store.register_backend(PythonRawFileBackend.KIND,
PythonRawFileBackend)
bar_store.register_backend(PythonRawFileBackend.KIND,
PythonRawFileBackend)
foo_data = {'foo': 'bar', 'bax': 'fox'}
bar_data = {'foo': 'bax', 'bax': 'foz'}
foo_store.put(foo_data, 'data')
bar_store.put(bar_data, 'data')
self.assertEqual(foo_store.get('data')[0], foo_data)
self.assertEqual(bar_store.get('data')[0], bar_data)
# -- files
foo_data = "some data"
file_like = BytesIO(foo_data.encode('utf-8'))
foo_store.put(file_like, 'myfile')
bar_data = "some other data"
file_like = BytesIO(bar_data.encode('utf-8'))
bar_store.put(file_like, 'myfile')
self.assertNotEqual(
foo_store.get('myfile').read(),
bar_store.get('myfile').read())
def test_put_dataframe_timestamp(self):
# create some dataframe
from datetime import datetime
df = pd.DataFrame({'a': list(range(1, 10)), 'b': list(range(1, 10))})
store = OmegaStore(prefix='')
# -- check default timestamp
now = datetime.utcnow()
store.put(df, 'mydata', append=False, timestamp=True)
df2 = store.get('mydata')
_created = df2['_created'].astype(datetime).unique()[0].to_pydatetime()
self.assertEqual(_created.replace(second=0, microsecond=0),
now.replace(second=0, microsecond=0))
# -- check custom timestamp column, default value
now = datetime.utcnow()
store.put(df, 'mydata', append=False, timestamp='CREATED')
df2 = store.get('mydata')
_created = df2['CREATED'].astype(datetime).unique()[0].to_pydatetime()
self.assertEqual(_created.replace(second=0, microsecond=0),
now.replace(second=0, microsecond=0))
# -- check custom timestamp column, value as tuple
now = datetime.utcnow() - timedelta(days=1)
store.put(df, 'mydata', append=False, timestamp=('CREATED', now))
df2 = store.get('mydata')
_created = df2['CREATED'].astype(datetime).unique()[0].to_pydatetime()
self.assertEqual(_created.replace(second=0, microsecond=0),
now.replace(second=0, microsecond=0))
# set a day in the past to avoid accidentally creating the current
# datetime in mongo
now = datetime.now() - timedelta(days=1)
store.put(df, 'mydata', timestamp=now, append=False)
df2 = store.get('mydata')
# compare the data
_created = df2['_created'].astype(datetime).unique()[0].to_pydatetime()
self.assertEqual(_created.replace(microsecond=0),
now.replace(microsecond=0))
def test_put_dataframe_timeseries(self):
# create some dataframe
tsidx = pd.date_range(pd.datetime(2016, 1, 1), pd.datetime(2016, 4, 1))
df = pd.DataFrame({
'a': list(range(0, len(tsidx))),
'b': list(range(0, len(tsidx)))
}, index=tsidx)
store = OmegaStore(prefix='')
store.put(df, 'mydata')
dfx = store.get('mydata')
assert_frame_equal(df, dfx)
idxs = list(store.collection('mydata').list_indexes())
idx_names = [dict(v).get('name') for v in idxs]
self.assertIn('asc__idx#0_0', idx_names)
def test_store_tz_datetime_dst(self):
""" test storing timezoned datetimes """
# 2019 11 03 02:00 is the end of US DST https://www.timeanddate.com/time/dst/2019.html
# pymongo will transform the object into a naive dt at UTC time at +3h (arguably incorrectly so)
# while pandas creates the Timestamp as UTC -4 (as the day starts at 00:00, not 02:00).
# On rendering back to a tz-aware datetime, this yields the wrong date (1 day eaerlier) because
# pandas applies -4 on converting from UTC to US/Eastern (correctly).
df = pd.DataFrame({
'y': pd.date_range('2019-11-01', periods=5, tz='US/Eastern', normalize=True)
})
store = OmegaStore()
store.put(df, 'test-date', append=False)
df2 = store.get('test-date')
# currently this fails, see @skip reason
testing.assert_frame_equal(df, df2)
def test_raw_files(self):
store = OmegaStore()
store.register_backend(PythonRawFileBackend.KIND, PythonRawFileBackend)
# test we can write from a file-like object
data = "some data"
file_like = BytesIO(data.encode('utf-8'))
store.put(file_like, 'myfile')
self.assertEqual(data.encode('utf-8'), store.get('myfile').read())
# test we can write from an actual file
data = "some other data"
file_like = BytesIO(data.encode('utf-8'))
with open('/tmp/testfile.txt', 'wb') as fout:
fout.write(file_like.read())
store.put('/tmp/testfile.txt', 'myfile')
self.assertEqual(data.encode('utf-8'), store.get('myfile').read())
def test_put_dataframe_multiindex(self):
# create some dataframe
store = OmegaStore(prefix='')
midx = pd.MultiIndex(levels=[[u'bar', u'baz', u'foo', u'qux'],
[u'one', u'two']],
labels=[[0, 0, 1, 1, 2, 2, 3, 3],
[0, 1, 0, 1, 0, 1, 0, 1]],
names=[u'first', u'second'])
df = pd.DataFrame({'x': range(0, len(midx))}, index=midx)
store.put(df, 'mydata')
dfx = store.get('mydata')
assert_frame_equal(df, dfx)
idxs = list(store.collection('mydata').list_indexes())
idx_names = [dict(v).get('name') for v in idxs]
self.assertIn('asc__idx#0_first__asc__idx#1_second', idx_names)
def test_put_model_with_prefix(self):
# create a test model
iris = load_iris()
X = iris.data
Y = iris.target
lr = LogisticRegression()
lr.fit(X, Y)
result = lr.predict(X)
# store it remote
store = OmegaStore(prefix='models/')
store.put(lr, 'foo')
# get it back, try predicting
lr2 = store.get('foo')
self.assertIsInstance(lr2, LogisticRegression)
result2 = lr2.predict(X)
self.assertTrue((result == result2).all())
def test_put_model(self):
# create a test model
iris = load_iris()
X = iris.data
Y = iris.target
lr = LogisticRegression(solver='liblinear', multi_class='auto')
lr.fit(X, Y)
result = lr.predict(X)
# store it remote
store = OmegaStore()
store.put(lr, 'models/foo')
# get it back, try predicting
lr2 = store.get('models/foo')
self.assertIsInstance(lr2, LogisticRegression)
result2 = lr2.predict(X)
self.assertTrue((result == result2).all())
def test_list_raw(self):
data = {'a': list(range(1, 10)), 'b': list(range(1, 10))}
df = pd.DataFrame(data)
store = OmegaStore()
meta = store.put(df, 'hdfdf', as_hdf=True)
# list with pattern
entries = store.list(pattern='hdf*', raw=True)
self.assertTrue(isinstance(entries[0], Metadata))
self.assertEqual('hdfdf', entries[0].name)
self.assertEqual(len(entries), 1)
# list with regexp
entries = store.list(regexp='hdf.*', raw=True)
self.assertTrue(isinstance(entries[0], Metadata))
self.assertEqual('hdfdf', entries[0].name)
self.assertEqual(len(entries), 1)
# list without pattern nor regexp
entries = store.list('hdfdf', kind=Metadata.PANDAS_HDF, raw=True)
self.assertTrue(isinstance(entries[0], Metadata))
self.assertEqual('hdfdf', entries[0].name)
self.assertEqual(len(entries), 1)
# subset kind
entries = store.list('hdfdf', raw=True, kind=Metadata.PANDAS_DFROWS)
self.assertEqual(len(entries), 0)
entries = store.list('hdfdf', raw=True, kind=Metadata.PANDAS_HDF)
self.assertEqual(len(entries), 1)
