diff --git a/tests/pipeline/test_blaze.py b/tests/pipeline/test_blaze.py index 0cbd8d84..202342e0 100644 --- a/tests/pipeline/test_blaze.py +++ b/tests/pipeline/test_blaze.py @@ -52,6 +52,19 @@ asset_infos = ( with_extra_sid = parameterized.expand(asset_infos) +def _utc_localize_index_level_0(df): + """``tz_localize`` the first level of a multiindexed dataframe to utc. + + Mutates df in place. + """ + idx = df.index + df.index = pd.MultiIndex.from_product( + (idx.levels[0].tz_localize('utc'), idx.levels[1]), + names=idx.names, + ) + return df + + class BlazeToPipelineTestCase(TestCase): @classmethod def setUpClass(cls): @@ -315,7 +328,11 @@ class BlazeToPipelineTestCase(TestCase): finder, ).run_pipeline(p, dates[0], dates[-1]) - assert_frame_equal(result, expected, check_dtype=False) + assert_frame_equal( + result, + _utc_localize_index_level_0(expected), + check_dtype=False, + ) def test_custom_query_time_tz(self): df = self.df.copy() @@ -345,7 +362,7 @@ class BlazeToPipelineTestCase(TestCase): expected = df.drop('asof_date', axis=1) expected['timestamp'] = expected['timestamp'].dt.normalize().astype( 'datetime64[ns]', - ) + ).dt.tz_localize('utc') expected.ix[3:5, 'timestamp'] += timedelta(days=1) expected.set_index(['timestamp', 'sid'], inplace=True) expected.index = pd.MultiIndex.from_product(( @@ -615,41 +632,26 @@ class BlazeToPipelineTestCase(TestCase): df['other'] = df.value + 1 fields = OrderedDict(self.macro_dshape.measure.fields) fields['other'] = fields['value'] - expr = bz.Data(df, name='expr', dshape=var * Record(fields)) - loader = BlazeLoader() - ds = from_blaze( - expr, - loader=loader, - no_deltas_rule=no_deltas_rules.ignore, - ) - p = Pipeline() - p.add(ds.value.latest, 'value') - p.add(ds.other.latest, 'other') - dates = self.dates asset_info = asset_infos[0][0] with tmp_asset_finder(equities=asset_info) as finder: - result = SimplePipelineEngine( - loader, - dates, + expected = pd.DataFrame( + np.array([[0, 1], + [1, 2], + [2, 3]]).repeat(3, axis=0), + index=pd.MultiIndex.from_product(( + df.timestamp, + finder.retrieve_all(asset_info.index), + )), + columns=('value', 'other'), + ).sort_index(axis=1) + self._test_id( + df, + var * Record(fields), + expected, finder, - ).run_pipeline(p, dates[0], dates[-1]) - - expected = pd.DataFrame( - np.array([[0, 1], - [1, 2], - [2, 3]]).repeat(3, axis=0), - index=pd.MultiIndex.from_product(( - df.timestamp, - finder.retrieve_all(asset_info.index), - )), - columns=('value', 'other'), - ).sort_index(axis=1) - assert_frame_equal( - result, - expected.sort_index(axis=1), - check_dtype=False, - ) + ('value', 'other'), + ) def test_id_take_last_in_group(self): T = pd.Timestamp @@ -804,7 +806,7 @@ class BlazeToPipelineTestCase(TestCase): assert_frame_equal( result, - expected_output, + _utc_localize_index_level_0(expected_output), check_dtype=False, ) diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 7b8b71a8..859adcec 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -174,6 +174,7 @@ from zipline.utils.input_validation import ( optionally, ) from zipline.utils.numpy_utils import repeat_last_axis +from zipline.utils.pandas_utils import sort_values from zipline.utils.preprocess import preprocess @@ -626,10 +627,10 @@ def overwrite_novel_deltas(baseline, deltas, dates): ) <= 1 novel_deltas = deltas.loc[novel_idx] non_novel_deltas = deltas.loc[~novel_idx] - return pd.concat( + return sort_values(pd.concat( (baseline, novel_deltas), ignore_index=True, - ).sort(TS_FIELD_NAME), non_novel_deltas + ), TS_FIELD_NAME), non_novel_deltas def overwrite_from_dates(asof, dense_dates, sparse_dates, asset_idx, value): @@ -910,13 +911,13 @@ class BlazeLoader(dict): This can return more data than needed. The in memory reindex will handle this. """ - return reduce( + return sort_values(reduce( partial(pd.merge, on=added_query_fields, how='outer'), ( odo(where(e, column), pd.DataFrame, **odo_kwargs) for column in columns ), - ).sort(TS_FIELD_NAME) # sort for the groupby later + ), TS_FIELD_NAME) # sort for the groupby later materialized_expr = collect_expr(expr) materialized_deltas = ( diff --git a/zipline/utils/pandas_utils.py b/zipline/utils/pandas_utils.py index e488ce84..2dd282e9 100644 --- a/zipline/utils/pandas_utils.py +++ b/zipline/utils/pandas_utils.py @@ -1,6 +1,7 @@ """ Utilities for working with pandas objects. """ +import pandas as pd def explode(df): @@ -10,3 +11,10 @@ def explode(df): (df.index, df.columns, df.values) """ return df.index, df.columns, df.values + + +try: + # pandas 0.16 compat + sort_values = pd.DataFrame.sort_values +except AttributeError: + sort_values = pd.DataFrame.sort