PERF: Optimize session close lookups in resample bar reader (#1749)

Optimize session close lookups in MinuteResampleSessionBarReader:

- Adds `session_closes_in_range` method (along with
  `session_opens_in_range`) to TradingCalendar to allow vectorized
  retrieval of all values in a range of sessions.
- Improves code path for resampling a single session's worth of data (as
  is the case when calling `get_value`), since we don't actually need to
  look up the close minute.

tests/calendars/test_trading_calendar.py

@@ -34,6 +34,7 @@ from zipline.errors import (
     InvalidCalendarName,
 )
 
+from zipline.testing.predicates import assert_equal
 from zipline.utils.calendars import(
     register_calendar,
     deregister_calendar,
@@ -683,6 +684,22 @@ class ExchangeCalendarTestBase(object):
             self.assertEqual(open_answer, found_open)
             self.assertEqual(close_answer, found_close)
 
+    def test_session_opens_in_range(self):
+        found_opens = self.calendar.session_opens_in_range(
+            self.answers.index[0],
+            self.answers.index[-1],
+        )
+
+        assert_equal(found_opens, self.answers['market_open'])
+
+    def test_session_closes_in_range(self):
+        found_closes = self.calendar.session_closes_in_range(
+            self.answers.index[0],
+            self.answers.index[-1],
+        )
+
+        assert_equal(found_closes, self.answers['market_close'])
+
     def test_daylight_savings(self):
         # 2004 daylight savings switches:
         # Sunday 2004-04-04 and Sunday 2004-10-31

zipline/data/resample.py

@@ -514,28 +514,49 @@ class MinuteResampleSessionBarReader(SessionBarReader):
         self._calendar = calendar
         self._minute_bar_reader = minute_bar_reader
 
-    def _get_resampled(self, columns, start_dt, end_dt, assets):
+    def _get_resampled(self, columns, start_session, end_session, assets):
+        range_open = self._calendar.session_open(start_session)
+        range_close = self._calendar.session_close(end_session)
+
         minute_data = self._minute_bar_reader.load_raw_arrays(
-            columns, start_dt, end_dt, assets)
-        dts = self._calendar.minutes_in_range(start_dt, end_dt).values
-        sessions = self._calendar.sessions_in_range(start_dt, end_dt)
-        m_closes = np.zeros(len(sessions), dtype=np.dtype('datetime64[ns]'))
-        for i, s in enumerate(sessions):
-            close = self._calendar.session_close(s)
-            m_closes[i] = close.value
-        m_locs = np.searchsorted(dts, m_closes)
+            columns,
+            range_open,
+            range_close,
+            assets,
+        )
+
+        # Get the index of the close minute for each session in the range.
+        # If the range contains only one session, the only close in the range
+        # is the last minute in the data. Otherwise, we need to get all the
+        # session closes and find their indices in the range of minutes.
+        if start_session == end_session:
+            close_ilocs = np.array([len(minute_data[0]) - 1], dtype=np.int64)
+        else:
+            minutes = self._calendar.minutes_in_range(
+                range_open,
+                range_close,
+            )
+            session_closes = self._calendar.session_closes_in_range(
+                start_session,
+                end_session,
+            )
+            close_ilocs = minutes.searchsorted(session_closes.values)
+
         results = []
-        shape = (len(sessions), len(assets))
+        shape = (len(close_ilocs), len(assets))
+
         for col in columns:
             if col != 'volume':
                 out = np.full(shape, np.nan)
             else:
                 out = np.zeros(shape, dtype=np.uint32)
             results.append(out)
+
         for i in range(len(assets)):
             for j, column in enumerate(columns):
                 data = minute_data[j][:, i]
-                minute_to_session(column, m_locs, data, results[j][:, i])
+                minute_to_session(column, close_ilocs, data, results[j][:, i])
+
         return results
 
     @property
@@ -543,19 +564,14 @@ class MinuteResampleSessionBarReader(SessionBarReader):
         return self._calendar
 
     def load_raw_arrays(self, columns, start_dt, end_dt, sids):
-        range_open, _ = self._calendar.open_and_close_for_session(
-            start_dt)
-        _, range_close = self._calendar.open_and_close_for_session(
-            end_dt)
-        return self._get_resampled(columns, range_open, range_close, sids)
+        return self._get_resampled(columns, start_dt, end_dt, sids)
 
     def get_value(self, sid, session, colname):
         # WARNING: This will need caching or other optimization if used in a
         # tight loop.
         # This was developed to complete interface, but has not been tuned
         # for real world use.
-        start, end = self._calendar.open_and_close_for_session(session)
-        return self._get_resampled([colname], start, end, [sid])[0][0][0]
+        return self._get_resampled([colname], session, session, [sid])[0][0][0]
 
     @lazyval
     def sessions(self):

zipline/utils/calendars/trading_calendar.py

@@ -660,6 +660,18 @@ class TradingCalendar(with_metaclass(ABCMeta)):
             'market_close'
         ].tz_localize('UTC')
 
+    def session_opens_in_range(self, start_session_label, end_session_label):
+        return self.schedule.loc[
+            start_session_label:end_session_label,
+            'market_open',
+        ].dt.tz_localize('UTC')
+
+    def session_closes_in_range(self, start_session_label, end_session_label):
+        return self.schedule.loc[
+            start_session_label:end_session_label,
+            'market_close',
+        ].dt.tz_localize('UTC')
+
     @property
     def all_sessions(self):
         return self.schedule.index