Merge pull request #979 from quantopian/allow-cols

ENH: Allow passing of numpy arrays to writer.

tests/data/test_minute_bars.py

@@ -17,7 +17,7 @@ import os
 
 from unittest import TestCase
 
-from numpy import nan
+from numpy import nan, array
 from numpy.testing import assert_almost_equal
 from pandas import (
     DataFrame,
@@ -446,3 +446,57 @@ class BcolzMinuteBarTestCase(TestCase):
         volume_price = self.reader.get_value(sid, minute, 'volume')
 
         self.assertEquals(100.0, volume_price)
+
+    def test_write_cols(self):
+        minute_0 = self.market_opens[self.test_calendar_start]
+        minute_1 = minute_0 + timedelta(minutes=1)
+        sid = 1
+        cols = {
+            'open': array([10.0, 11.0]),
+            'high': array([20.0, 21.0]),
+            'low': array([30.0, 31.0]),
+            'close': array([40.0, 41.0]),
+            'volume': array([50.0, 51.0])
+        }
+        dts = array([minute_0, minute_1], dtype='datetime64[s]')
+        self.writer.write_cols(sid, dts, cols)
+
+        open_price = self.reader.get_value(sid, minute_0, 'open')
+
+        self.assertEquals(10.0, open_price)
+
+        high_price = self.reader.get_value(sid, minute_0, 'high')
+
+        self.assertEquals(20.0, high_price)
+
+        low_price = self.reader.get_value(sid, minute_0, 'low')
+
+        self.assertEquals(30.0, low_price)
+
+        close_price = self.reader.get_value(sid, minute_0, 'close')
+
+        self.assertEquals(40.0, close_price)
+
+        volume_price = self.reader.get_value(sid, minute_0, 'volume')
+
+        self.assertEquals(50.0, volume_price)
+
+        open_price = self.reader.get_value(sid, minute_1, 'open')
+
+        self.assertEquals(11.0, open_price)
+
+        high_price = self.reader.get_value(sid, minute_1, 'high')
+
+        self.assertEquals(21.0, high_price)
+
+        low_price = self.reader.get_value(sid, minute_1, 'low')
+
+        self.assertEquals(31.0, low_price)
+
+        close_price = self.reader.get_value(sid, minute_1, 'close')
+
+        self.assertEquals(41.0, close_price)
+
+        volume_price = self.reader.get_value(sid, minute_1, 'volume')
+
+        self.assertEquals(51.0, volume_price)

zipline/data/minute_bars.py

@@ -20,7 +20,6 @@ from os.path import join
 import json
 import os
 import pandas as pd
-from pandas.core.datetools import normalize_date
 
 US_EQUITIES_MINUTES_PER_DAY = 390
 
@@ -403,8 +402,6 @@ class BcolzMinuteBarWriter(object):
         -----------
         sid : int
             The asset identifer for the data being written.
-        days : pd.DatetimeIndex
-            The days for which to write data from the given df.
         df : pd.DataFrame
             DataFrame of market data with the following characteristics.
             columns : ('open', 'high', 'low', 'close', 'volume')
@@ -415,11 +412,49 @@ class BcolzMinuteBarWriter(object):
                 volume : float64|int64
             index : DatetimeIndex of market minutes.
         """
+        cols = {
+            'open': df.open.values,
+            'high': df.high.values,
+            'low': df.low.values,
+            'close': df.close.values,
+            'volume': df.volume.values,
+        }
+        dts = df.index.values
+        self.write_cols(sid, dts, cols)
+
+    def write_cols(self, sid, dts, cols):
+        """
+        Write the OHLCV data for the given sid.
+
+        If there is no bcolz ctable yet created for the sid, create it.
+
+        If the length of the bcolz ctable is not exactly to the date before
+        the first day provided, fill the ctable with 0s up to that date.
+
+        Writes in blocks of the size of the days times minutes per day.
+
+        Parameters:
+        -----------
+        sid : int
+            The asset identifer for the data being written.
+        dts : datetime64 array
+            The dts corresponding to values in cols.
+        cols : dict of str -> np.array
+            dict of market data with the following characteristics.
+            keys are ('open', 'high', 'low', 'close', 'volume')
+                open : float64
+                high : float64
+                low  : float64
+                close : float64
+                volume : float64|int64
+        """
         table = self._ensure_ctable(sid)
 
         tds = self._trading_days
-        input_first_day = normalize_date(df.index[0])
-        input_last_day = normalize_date(df.index[-1])
+        input_first_day = pd.Timestamp(dts[0].astype('datetime64[D]'),
+                                       tz='UTC')
+        input_last_day = pd.Timestamp(dts[-1].astype('datetime64[D]'),
+                                      tz='UTC')
 
         last_date = self.last_date_in_output_for_sid(sid)
 
@@ -449,15 +484,15 @@ class BcolzMinuteBarWriter(object):
         vol_col = np.zeros(minutes_count, dtype=np.uint32)
 
         dt_ixs = np.searchsorted(all_minutes_in_window.values,
-                                 df.index.values)
+                                 dts.astype('datetime64[ns]'))
 
         ohlc_ratio = self._ohlc_ratio
-        open_col[dt_ixs] = (df.open.values * ohlc_ratio).astype(np.uint32)
-        high_col[dt_ixs] = (df.high.values * ohlc_ratio).astype(np.uint32)
-        low_col[dt_ixs] = (df.low.values * ohlc_ratio).astype(np.uint32)
-        close_col[dt_ixs] = (df.close.values * ohlc_ratio).astype(
+        open_col[dt_ixs] = (cols['open'] * ohlc_ratio).astype(np.uint32)
+        high_col[dt_ixs] = (cols['high'] * ohlc_ratio).astype(np.uint32)
+        low_col[dt_ixs] = (cols['low'] * ohlc_ratio).astype(np.uint32)
+        close_col[dt_ixs] = (cols['close'] * ohlc_ratio).astype(
             np.uint32)
-        vol_col[dt_ixs] = df.volume.values.astype(np.uint32)
+        vol_col[dt_ixs] = cols['volume'].astype(np.uint32)
 
         table.append([
             open_col,