MAINT: move some blaze utilities into a shared module

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
+    )

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
         )