Adjustments to Component Dfs (#1620)

* ENH: SQLiteAdjustmentReader can return DF versions of tables.

tests/pipeline/test_us_equity_pricing_loader.py

@@ -15,6 +15,7 @@
 """
 Tests for USEquityPricingLoader and related classes.
 """
+from nose_parameterized import parameterized
 from numpy import (
     arange,
     datetime64,
@@ -32,6 +33,7 @@ from pandas import (
     Int64Index,
     Timestamp,
 )
+from pandas.util.testing import assert_frame_equal
 from toolz.curried.operator import getitem
 
 from zipline.lib.adjustment import Float64Multiply
@@ -395,6 +397,57 @@ class USEquityPricingLoaderTestCase(WithAdjustmentReader,
                 self.assertEqual(adj.last_col, expected.last_col)
                 assert_allclose(adj.value, expected.value)
 
+    @parameterized([(True,), (False,)])
+    def test_load_adjustments_to_df(self, convert_dts):
+        reader = self.adjustment_reader
+        adjustment_dfs = reader.unpack_db_to_component_dfs(
+            convert_dates=convert_dts
+        )
+
+        name_and_raw = (
+            ('splits', SPLITS),
+            ('mergers', MERGERS),
+            ('dividends', DIVIDENDS_EXPECTED)
+        )
+
+        def create_expected_table(df, name):
+            expected_df = df.copy()
+
+            if convert_dts:
+                for colname in reader._datetime_int_cols[name]:
+                    expected_df[colname] = expected_df[colname].astype(
+                        'datetime64[s]'
+                    )
+
+            return expected_df
+
+        def create_expected_div_table(df, name):
+            expected_df = df.copy()
+
+            if not convert_dts:
+                for colname in reader._datetime_int_cols[name]:
+                    expected_df[colname] = expected_df[colname].astype(
+                        'datetime64[s]'
+                    ).astype(int)
+
+            return expected_df
+
+        for action_name, raw_tbl in name_and_raw:
+
+            exp = create_expected_table(raw_tbl, action_name)
+            assert_frame_equal(
+                adjustment_dfs[action_name],
+                exp
+            )
+
+        # DIVIDENDS is in the opposite form from the rest of the dataframes, so
+        # needs to be converted separately.
+        div_name = 'dividend_payouts'
+        assert_frame_equal(
+            adjustment_dfs[div_name],
+            create_expected_div_table(DIVIDENDS, div_name)
+        )
+
     def test_read_no_adjustments(self):
         adjustment_reader = NullAdjustmentReader()
         columns = [USEquityPricing.close, USEquityPricing.volume]

zipline/data/us_equity_pricing.py

@@ -36,18 +36,19 @@ from numpy import (
     uint32,
 )
 from pandas import (
-    isnull,
     DataFrame,
-    read_csv,
-    Timestamp,
+    DatetimeIndex,
+    isnull,
     NaT,
-    DatetimeIndex
+    read_csv,
+    read_sql,
+    Timestamp,
 )
 from pandas.tslib import iNaT
 from six import (
     iteritems,
-    viewkeys,
     string_types,
+    viewkeys,
 )
 
 from zipline.data.session_bars import SessionBarReader
@@ -60,8 +61,8 @@ from zipline.utils.calendars import get_calendar
 from zipline.utils.functional import apply
 from zipline.utils.preprocess import call
 from zipline.utils.input_validation import (
-    preprocess,
     expect_element,
+    preprocess,
     verify_indices_all_unique,
 )
 from zipline.utils.sqlite_utils import group_into_chunks, coerce_string_to_conn
@@ -1250,6 +1251,18 @@ class SQLiteAdjustmentReader(object):
     def __init__(self, conn):
         self.conn = conn
 
+        # Given the tables in the adjustments.db file, dict which knows which
+        # col names contain dates that have been coerced into ints.
+        self._datetime_int_cols = {
+            'dividend_payouts': ('declared_date', 'ex_date', 'pay_date',
+                                 'record_date'),
+            'dividends': ('effective_date',),
+            'mergers': ('effective_date',),
+            'splits': ('effective_date',),
+            'stock_dividend_payouts': ('declared_date', 'ex_date', 'pay_date',
+                                       'record_date')
+        }
+
     def load_adjustments(self, columns, dates, assets):
         return load_adjustments_from_sqlite(
             self.conn,
@@ -1316,3 +1329,46 @@ class SQLiteAdjustmentReader(object):
         c.close()
 
         return stock_divs
+
+    def unpack_db_to_component_dfs(self, convert_dates=False):
+        """Returns the set of known tables in the adjustments file in DataFrame
+        form.
+
+        Parameters
+        ----------
+        convert_dates : bool, optional
+            By default, dates are returned in seconds since EPOCH. If
+            convert_dates is True, all ints in date columns will be converted
+            to datetimes.
+
+        Returns
+        -------
+        dfs : dict{str->DataFrame}
+            Dictionary which maps table name to the corresponding DataFrame
+            version of the table, where all date columns have been coerced back
+            from int to datetime.
+        """
+
+        def _get_df_from_table(table_name, date_cols):
+
+            kwargs = (
+                {'parse_dates': {col: 's' for col in date_cols}}
+                if convert_dates
+                else {}
+            )
+
+            # Dates are stored in second resolution as ints in adj.db tables.
+            return read_sql(
+                'select * from "{}"'.format(table_name),
+                self.conn,
+                index_col='index',
+                **kwargs
+            ).rename_axis(None)
+
+        return {
+            t_name: _get_df_from_table(
+                t_name,
+                date_cols
+            )
+            for t_name, date_cols in self._datetime_int_cols.items()
+        }