From b037a06576d4918b3546d6082aa11870688de7d8 Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Tue, 5 Jan 2016 14:04:13 -0500 Subject: [PATCH] MAINT: move some blaze utilities into a shared module --- zipline/pipeline/loaders/blaze/core.py | 85 +++++++++++++++++++++- zipline/pipeline/loaders/blaze/earnings.py | 73 ++++--------------- 2 files changed, 98 insertions(+), 60 deletions(-) diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 8db633a6..556eb3f8 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -154,7 +154,7 @@ from toolz import ( memoize, ) import toolz.curried.operator as op -from six import with_metaclass, PY2, itervalues +from six import with_metaclass, PY2, itervalues, iteritems from zipline.pipeline.data.dataset import DataSet, Column @@ -903,3 +903,86 @@ class BlazeLoader(dict): ) global_loader = BlazeLoader.global_instance() + + +def bind_expression_to_resources(expr, resources): + """ + Bind a Blaze expression to resources. + + Parameters + ---------- + expr : bz.Expr + The expression to which we want to bind resources. + resources : dict[bz.Symbol -> any] + Mapping from the atomic terms of ``expr`` to actual data resources. + + Returns + ------- + bound_expr : bz.Expr + ``expr`` with bound resources. + """ + # bind the resources into the expression + if resources is None: + resources = {} + + # _subs stands for substitute. It's not actually private, blaze just + # prefixes symbol-manipulation methods with underscores to prevent + # collisions with data column names. + return expr._subs({ + k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources) + }) + + +def ffill_query_in_range(expr, + lower, + upper, + odo_kwargs=None, + ts_field=TS_FIELD_NAME, + sid_field=SID_FIELD_NAME): + """Query a blaze expression in a given time range properly forward filling + from values that fall before the lower date. + + Parameters + ---------- + expr : Expr + Bound blaze expression. + lower : datetime + The lower date to query for. + upper : datetime + The upper date to query for. + odo_kwargs : dict, optional + The extra keyword arguments to pass to ``odo``. + ts_field : str, optional + The name of the timestamp field in the given blaze expression. + sid_field : str, optional + The name of the sid field in the given blaze expression. + + Returns + ------- + raw : pd.DataFrame + A strict dataframe for the data in the given date range. This may + start before the requested start date if a value is needed to ffill. + """ + odo_kwargs = odo_kwargs or {} + filtered = expr[expr[ts_field] <= lower] + computed_lower = odo( + bz.by( + filtered[sid_field], + timestamp=filtered[ts_field].max(), + ).timestamp.min(), + pd.Timestamp, + **odo_kwargs + ) + if pd.isnull(computed_lower): + # If there is no lower date, just query for data in the date + # range. It must all be null anyways. + computed_lower = lower + + return odo( + expr[ + (expr[ts_field] >= computed_lower) & + (expr[ts_field] <= upper) + ], + pd.DataFrame, + **odo_kwargs + ) diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 6b5c1695..2c4f8dfb 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -1,11 +1,13 @@ -import blaze as bz from datashape import istabular -from odo import odo import pandas as pd -from six import iteritems from toolz import valmap -from .core import TS_FIELD_NAME, SID_FIELD_NAME +from .core import ( + TS_FIELD_NAME, + SID_FIELD_NAME, + bind_expression_to_resources, + ffill_query_in_range, +) from zipline.pipeline.data import EarningsCalendar from zipline.pipeline.loaders.base import PipelineLoader from zipline.pipeline.loaders.earnings import EarningsCalendarLoader @@ -14,34 +16,6 @@ from zipline.pipeline.loaders.earnings import EarningsCalendarLoader ANNOUNCEMENT_FIELD_NAME = 'announcement_date' -def bind_expression_to_resources(expr, resources): - """ - Bind a Blaze expression to resources. - - Parameters - ---------- - expr : bz.Expr - The expression to which we want to bind resources. - resources : dict[bz.Symbol -> any] - Mapping from the atomic terms of ``expr`` to actual data resources. - - Returns - ------- - bound_expr : bz.Expr - ``expr`` with bound resources. - """ - # bind the resources into the expression - if resources is None: - resources = {} - - # _subs stands for substitute. It's not actually private, blaze just - # prefixes symbol-manipulation methods with underscores to prevent - # collisions with data column names. - return expr._subs({ - k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources) - }) - - class BlazeEarningsCalendarLoader(PipelineLoader): """A pipeline loader for the ``EarningsCalendar`` dataset that loads data from a blaze expression. @@ -61,8 +35,8 @@ class BlazeEarningsCalendarLoader(PipelineLoader): Dim * {{ {SID_FIELD_NAME}: int64, - {TS_FIELD_NAME}: datetime64, - {ANNOUNCEMENT_FIELD_NAME}: datetime64, + {TS_FIELD_NAME}: datetime, + {ANNOUNCEMENT_FIELD_NAME}: ?datetime, }} Where each row of the table is a record including the sid to identify the @@ -87,7 +61,6 @@ class BlazeEarningsCalendarLoader(PipelineLoader): def __init__(self, expr, resources=None, - compute_kwargs=None, odo_kwargs=None, dataset=EarningsCalendar): dshape = expr.dshape @@ -106,33 +79,15 @@ class BlazeEarningsCalendarLoader(PipelineLoader): self._dataset = dataset def load_adjusted_array(self, columns, dates, assets, mask): - expr = self._expr - filtered = expr[expr[TS_FIELD_NAME] <= dates[0]] - lower = odo( - bz.by( - filtered[SID_FIELD_NAME], - timestamp=filtered[TS_FIELD_NAME].max(), - ).timestamp.min(), - pd.Timestamp, - **self._odo_kwargs + raw = ffill_query_in_range( + self._expr, + dates[0], + dates[-1], + self._odo_kwargs, ) - if pd.isnull(lower): - # If there is no lower date, just query for data in the date - # range. It must all be null anyways. - lower = dates[0] - - raw = odo( - expr[ - (expr[TS_FIELD_NAME] >= lower) & - (expr[TS_FIELD_NAME] <= dates[-1]) - ], - pd.DataFrame, - **self._odo_kwargs - ) - sids = raw.loc[:, SID_FIELD_NAME] raw.drop( - sids[~(sids.isin(assets) | sids.notnull())].index, + sids[~sids.isin(assets)].index, inplace=True )