Merge pull request #1492 from quantopian/cache-minute-to-session-lookups

PERF: Save up to 80% of the calls to minute_to_session_label

zipline/utils/calendars/trading_calendar.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABCMeta, abstractproperty
+from lru import LRU
 
 from pandas.tseries.holiday import AbstractHolidayCalendar
 from six import with_metaclass
@@ -95,6 +96,12 @@ class TradingCalendar(with_metaclass(ABCMeta)):
             dtype='datetime64[ns]',
         )
 
+        # Simple cache to avoid recalculating the same minute -> session in
+        # "next" mode. Analysis of current zipline code paths show that
+        # `minute_to_session_label` is often called consecutively with the same
+        # inputs.
+        self._minute_to_session_label_cache = LRU(1)
+
         self.market_opens_nanos = self.schedule.market_open.values.\
             astype(np.int64)
 
@@ -690,11 +697,19 @@ class TradingCalendar(with_metaclass(ABCMeta)):
         pd.Timestamp (midnight UTC)
             The label of the containing session.
         """
+        if direction == "next":
+            try:
+                return self._minute_to_session_label_cache[dt]
+            except KeyError:
+                pass
 
         idx = searchsorted(self.market_closes_nanos, dt)
         current_or_next_session = self.schedule.index[idx]
+        self._minute_to_session_label_cache[dt] = current_or_next_session
 
-        if direction == "previous":
+        if direction == "next":
+            return current_or_next_session
+        elif direction == "previous":
             if not is_open(self.market_opens_nanos, self.market_closes_nanos,
                            dt):
                 # if the exchange is closed, use the previous session
@@ -704,7 +719,7 @@ class TradingCalendar(with_metaclass(ABCMeta)):
                            dt):
                 # if the exchange is closed, blow up
                 raise ValueError("The given dt is not an exchange minute!")
-        elif direction != "next":
+        else:
             # invalid direction
             raise ValueError("Invalid direction parameter: "
                              "{0}".format(direction))