diff --git a/tests/pipeline/test_blaze.py b/tests/pipeline/test_blaze.py index 11d741b2..a8b57e8c 100644 --- a/tests/pipeline/test_blaze.py +++ b/tests/pipeline/test_blaze.py @@ -328,8 +328,8 @@ class BlazeToPipelineTestCase(TestCase): df['timestamp'] = ( pd.DatetimeIndex(df['timestamp'], tz='EST') + timedelta(hours=8, minutes=44) - ).tz_convert('utc') - df.ix[3:5, 'timestamp'] = pd.Timestamp('2014-01-01 13:45', tz='utc') + ).tz_convert('utc').tz_localize(None) + df.ix[3:5, 'timestamp'] = pd.Timestamp('2014-01-01 13:45') expr = bz.Data(df, name='expr', dshape=self.dshape) loader = BlazeLoader(data_query_time=time(8, 45), data_query_tz='EST') ds = from_blaze( diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 75e365da..fb79db8d 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -127,7 +127,7 @@ from __future__ import division, absolute_import from abc import ABCMeta, abstractproperty from collections import namedtuple, defaultdict from copy import copy -from functools import partial +from functools import partial, reduce from itertools import count import warnings from weakref import WeakKeyDictionary @@ -188,7 +188,7 @@ traversable_nodes = ( bz.expr.Label, ) is_invalid_deltas_node = complement(flip(isinstance, valid_deltas_node_types)) -getname = op.attrgetter('__name__') +get__name__ = op.attrgetter('__name__') class _ExprRepr(object): @@ -523,8 +523,10 @@ def from_blaze(expr, raise TypeError( 'expression with deltas may only contain (%s) nodes,' " found: %s" % ( - ', '.join(map(getname, valid_deltas_node_types)), - ', '.join(set(map(compose(getname, type), invalid_nodes))), + ', '.join(map(get__name__, valid_deltas_node_types)), + ', '.join( + set(map(compose(get__name__, type), invalid_nodes)), + ), ), ) @@ -602,7 +604,7 @@ def from_blaze(expr, getdataset = op.attrgetter('dataset') -dataset_name = op.attrgetter('name') +getname = op.attrgetter('name') def overwrite_novel_deltas(baseline, deltas, dates): @@ -778,8 +780,10 @@ class BlazeLoader(dict): Parameters ---------- - colmap : mapping[BoundColumn -> tuple[Expr, Expr, any]], optional - The initial column mapping to use. + dsmap : mapping, optional + An initial mapping of datasets to ``ExprData`` objects. + NOTE: Further mutations to this map will not be reflected by this + object. data_query_time : time, optional The time to use for the data query cutoff. data_query_tz : tzinfo or str @@ -787,11 +791,10 @@ class BlazeLoader(dict): """ @preprocess(data_query_tz=optionally(ensure_timezone)) def __init__(self, - colmap=None, + dsmap=None, data_query_time=None, data_query_tz=None): - self.update(colmap or {}) - + self.update(dsmap or {}) check_data_query_args(data_query_time, data_query_tz) self._data_query_time = data_query_time self._data_query_tz = data_query_tz @@ -826,8 +829,7 @@ class BlazeLoader(dict): expr, deltas, resources = self[dataset] have_sids = SID_FIELD_NAME in expr.fields assets = list(map(int, assets)) # coerce from numpy.int64 - fields = list(map(dataset_name, columns)) - query_fields = fields + [AD_FIELD_NAME, TS_FIELD_NAME] + ( + added_query_fields = [AD_FIELD_NAME, TS_FIELD_NAME] + ( [SID_FIELD_NAME] if have_sids else [] ) @@ -840,9 +842,47 @@ class BlazeLoader(dict): data_query_tz, ) - def where(e): + def where(e, column): """Create the query to run against the resources. + Parameters + ---------- + e : Expr + The baseline or deltas expression. + column : BoundColumn + The column to query for. + + Returns + ------- + q : Expr + The query to run for the given column. + """ + colname = column.name + filtered = e[e[colname].notnull() & (e[TS_FIELD_NAME] <= lower_dt)] + lower = filtered.timestamp.max() + + if have_sids: + # If we have sids, then we need to take the earliest of the + # greatest date that has a non-null value by sid. + lower = bz.by( + filtered[SID_FIELD_NAME], + timestamp=lower, + ).timestamp.min() + + lower = odo(lower, pd.Timestamp) + if lower is pd.NaT: + # If there is no lower date, just query for data in he date + # range. It must all be null anyways. + lower = lower_dt + + return e[ + (e[TS_FIELD_NAME] >= lower) & + (e[TS_FIELD_NAME] <= upper_dt) + ][added_query_fields + [colname]] + + def collect_expr(e, _kwargs={'d': resources} if resources else {}): + """Execute and merge all of the per-column subqueries. + Parameters ---------- e : Expr @@ -850,28 +890,29 @@ class BlazeLoader(dict): Returns ------- - q : Expr - The query to run. + result : pd.DataFrame + The resulting dataframe. + + Notes + ----- + This can return more data than needed. The in memory reindex will + handle this. """ - ts = e[TS_FIELD_NAME] - # Hack to get the lower bound to query: - # This must be strictly executed because the data for `ts` will - # be removed from scope too early otherwise. - lower = odo(ts[ts <= lower_dt].max(), pd.Timestamp) - selection = ts <= upper_dt - if have_sids: - selection &= e[SID_FIELD_NAME].isin(assets) - if lower is not pd.NaT: - selection &= ts >= lower + return reduce( + partial(pd.merge, on=added_query_fields, how='outer'), + ( + odo(where(e, column), pd.DataFrame, **_kwargs) + for column in columns + ), + ) - return e[selection][query_fields] - - extra_kwargs = {'d': resources} if resources else {} - materialized_expr = odo(where(expr), pd.DataFrame, **extra_kwargs) + materialized_expr = collect_expr(expr) materialized_deltas = ( - odo(where(deltas), pd.DataFrame, **extra_kwargs) + collect_expr(deltas) if deltas is not None else - pd.DataFrame(columns=query_fields) + pd.DataFrame( + columns=added_query_fields + list(map(getname, columns)), + ) ) if data_query_time is not None: @@ -907,10 +948,11 @@ class BlazeLoader(dict): sparse_deltas = non_novel_deltas.set_index( [TS_FIELD_NAME, SID_FIELD_NAME], ).unstack() - - dense_output = sparse_output.reindex(dates, method='ffill') - cols = dense_output.columns - dense_output = dense_output.reindex( + cols = sparse_output.columns + dense_output = sparse_output.groupby( + dates[dates.searchsorted(sparse_output.index)], + ).last().reindex( + index=dates, columns=pd.MultiIndex.from_product( (cols.levels[0], assets), names=cols.names, @@ -933,10 +975,14 @@ class BlazeLoader(dict): partial(repeat_last_axis, count=len(assets)), ) sparse_output = sparse_output.set_index(TS_FIELD_NAME) - dense_output = sparse_output.reindex(dates, method='ffill') + dense_output = sparse_output.groupby( + dates[dates.searchsorted(sparse_output.index)], + ).last().reindex(dates) sparse_deltas = non_novel_deltas.set_index(TS_FIELD_NAME) adjustments_from_deltas = adjustments_from_deltas_no_sids + dense_output.ffill(inplace=True) + for column_idx, column in enumerate(columns): column_name = column.name yield column, AdjustedArray(