ENH: Add function for running pipelines in chunks

tests/pipeline/test_run_chunked_pipeline.py

@@ -0,0 +1,37 @@
+from zipline.pipeline import Pipeline, run_chunked_pipeline
+from zipline.pipeline.data import USEquityPricing
+from zipline.pipeline.factors import Returns
+from zipline.testing import ZiplineTestCase
+from zipline.testing.fixtures import WithEquityPricingPipelineEngine
+
+
+class ChunkedPipelineTestCase(WithEquityPricingPipelineEngine,
+                              ZiplineTestCase):
+
+    def test_run_chunked_pipeline(self):
+        """
+        Test that running a pipeline in chunks produces the same result as if
+        it were run all at once
+        """
+        pipe = Pipeline(
+            columns={
+                'close': USEquityPricing.close.latest,
+                'returns': Returns(window_length=2),
+            },
+        )
+        sessions = self.nyse_calendar.all_sessions
+        start_date = sessions[sessions.get_loc(self.START_DATE) + 2]
+
+        pipeline_result = self.pipeline_engine.run_pipeline(
+            pipe,
+            start_date=start_date,
+            end_date=self.END_DATE,
+        )
+        chunked_result = run_chunked_pipeline(
+            engine=self.pipeline_engine,
+            pipeline=pipe,
+            start_date=start_date,
+            end_date=self.END_DATE,
+            chunksize=22
+        )
+        self.assertTrue(chunked_result.equals(pipeline_result))

tests/utils/test_date_utils.py

@@ -1,5 +1,4 @@
 from pandas import Timestamp
-
 from nose_parameterized import parameterized
 
 from zipline.testing import ZiplineTestCase

zipline/pipeline/__init__.py

@@ -1,5 +1,8 @@
 from __future__ import print_function
 from zipline.assets import AssetFinder
+from zipline.utils.calendars import get_calendar
+from zipline.utils.date_utils import compute_date_range_chunks
+from zipline.utils.pandas_utils import categorical_df_concat
 
 from .classifiers import Classifier, CustomClassifier
 from .engine import SimplePipelineEngine
@@ -47,6 +50,39 @@ def engine_from_files(daily_bar_path,
     )
 
 
+def run_chunked_pipeline(engine, pipeline, start_date, end_date, chunksize):
+    """Run a pipeline to collect the results.
+
+    Parameters
+    ----------
+    engine : Engine
+        The pipeline engine.
+    pipeline : Pipeline
+        The pipeline to run.
+    start_date : pd.Timestamp
+        The start date to run the pipeline for.
+    end_date : pd.Timestamp
+        The end date to run the pipeline for.
+    chunksize : int or None
+        The number of days to execute at a time. If this is None, all the days
+        will be run at once.
+
+    Returns
+    -------
+    results : pd.DataFrame
+        The results for each output term in the pipeline.
+    """
+    ranges = compute_date_range_chunks(
+        get_calendar('NYSE'),
+        start_date,
+        end_date,
+        chunksize,
+    )
+    chunks = [engine.run_pipeline(pipeline, s, e) for s, e in ranges]
+
+    return categorical_df_concat(chunks, inplace=True)
+
+
 __all__ = (
     'Classifier',
     'CustomFactor',
@@ -58,6 +94,7 @@ __all__ = (
     'Filter',
     'Pipeline',
     'SimplePipelineEngine',
+    'run_chunked_pipeline',
     'Term',
     'TermGraph',
 )

zipline/utils/date_utils.py

@@ -1,11 +1,15 @@
-def roll_dates_to_previous_session(calendar, *dates):
+from toolz import partition_all
+
+
+def roll_dates_to_previous_session(sessions, *dates):
     """
-    Roll ``dates`` to the next session of ``calendar``.
+    Roll `dates` to the last session of `calendar`. Return input date if it
+    is a valid session.
 
     Parameters
     ----------
-    calendar : zipline.utils.calendars.trading_calendar.TradingCalendar
-        The calendar to use as a reference.
+    sessions : pandas.tseries.index.DatetimeIndex
+        The list of valid session dates.
     *dates : pd.Timestamp
         The dates for which the last trading date is needed.
 
@@ -15,7 +19,47 @@ def roll_dates_to_previous_session(calendar, *dates):
         The last trading date of the input dates, inclusive.
 
     """
-    all_sessions = calendar.all_sessions
+    # Find the previous index value if there is no exact match.
+    locs = [sessions.get_loc(dt, method='ffill') for dt in dates]
+    return sessions[locs].tolist()
 
-    locs = [all_sessions.get_loc(dt, method='ffill') for dt in dates]
-    return all_sessions[locs]
+
+def compute_date_range_chunks(sessions, start_date, end_date, chunksize):
+    """Compute the start and end dates to run a pipeline for.
+
+    Parameters
+    ----------
+    sessions : DatetimeIndex
+        The available dates.
+    start_date : pd.Timestamp
+        The first date in the pipeline.
+    end_date : pd.Timestamp
+        The last date in the pipeline.
+    chunksize : int or None
+        The size of the chunks to run. Setting this to None returns one chunk.
+
+    Returns
+    -------
+    ranges : iterable[(np.datetime64, np.datetime64)]
+        A sequence of start and end dates to run the pipeline for.
+    """
+    if start_date not in sessions:
+        raise KeyError("Start date %s is not found in calendar." %
+                       (start_date.strftime("%Y-%m-%d"),))
+    if end_date not in sessions:
+        raise KeyError("End date %s is not found in calendar." %
+                       (end_date.strftime("%Y-%m-%d"),))
+    if end_date < start_date:
+        raise ValueError("End date %s cannot precede start date %s." %
+                         (end_date.strftime("%Y-%m-%d"),
+                          start_date.strftime("%Y-%m-%d")))
+
+    if chunksize is None:
+        return [(start_date, end_date)]
+
+    start_ix, end_ix = sessions.slice_locs(start_date, end_date)
+    return (
+        (r[0], r[-1]) for r in partition_all(
+            chunksize, sessions[start_ix:end_ix]
+        )
+    )

zipline/utils/pandas_utils.py

@@ -258,7 +258,7 @@ def categorical_df_concat(df_list, inplace=False):
         new_categories = sorted(
             set().union(
                 *(frame[col].cat.categories for frame in df_list)
-            )
+            ) - {None}
         )
 
         for df in df_list: