Merge pull request pandas-dev#208 from manahl/chunkstore-dr

Date Chunker support for DatetimeIndexes

Author: bmoscon

arctic/chunkstore/chunkstore.py

@@ -94,7 +94,7 @@ def delete(self, symbol, chunk_range=None):
         chunk_range: range object
             a date range to delete
         """
-        if chunk_range:
+        if chunk_range is not None:
             # read out chunks that fall within the range and filter out
             # data within the range
             df = self.read(symbol, chunk_range=chunk_range, filter_data=False)
@@ -166,7 +166,8 @@ def read(self, symbol, chunk_range=None, columns=None, filter_data=True):
             the symbol to retrieve
         chunk_range: object
             corresponding range object for the specified chunker (for
-            DateChunker it is a DateRange object)
+            DateChunker it is a DateRange object or a DatetimeIndex, 
+            as returned by pandas.date_range
         columns: list of str
             subset of columns to read back (index will always be included, if
             one exists)
@@ -186,7 +187,7 @@ def read(self, symbol, chunk_range=None, columns=None, filter_data=True):
         spec = {SYMBOL: symbol,
                 }
 
-        if chunk_range:
+        if chunk_range is not None:
             spec.update(self.chunker.to_mongo(chunk_range))
 
         segments = []
@@ -284,7 +285,7 @@ def __update(self, symbol, item, combine_method=None, chunk_range=None):
         if sym[TYPE] == 'dataframe' and not isinstance(item, DataFrame):
             raise Exception("Cannot combine DataFrame and Series")
 
-        if chunk_range:
+        if chunk_range is not None:
             self.delete(symbol, chunk_range)
             sym = self._get_symbol_info(symbol)
 
@@ -365,7 +366,7 @@ def update(self, symbol, item, chunk_range=None):
             original data.
         """
 
-        if chunk_range:
+        if chunk_range is not None:
             if self.chunker.filter(item, chunk_range).empty:
                 raise Exception('Range must be inclusive of data')
             self.__update(symbol, item, combine_method=self.__concat, chunk_range=chunk_range)

arctic/chunkstore/date_chunker.py

@@ -61,6 +61,8 @@ def to_mongo(self, range_obj):
         -------
         string
         """
+        if isinstance(range_obj, pd.DatetimeIndex):
+            range_obj = DateRange(range_obj.min(), range_obj.max())
         if range_obj.start and range_obj.end:
             return {'$and': [{START: {'$lte': range_obj.end}}, {END: {'$gte': range_obj.start}}]}
         elif range_obj.start:
@@ -83,6 +85,8 @@ def filter(self, data, range_obj):
         -------
         data, filtered by range_obj
         """
+        if isinstance(range_obj, pd.DatetimeIndex):
+            range_obj = DateRange(range_obj.min(), range_obj.max())
         if 'date' in data.index.names:
             return data[range_obj.start:range_obj.end]
         elif 'date' in data.columns:
@@ -98,6 +102,8 @@ def exclude(self, data, range_obj):
         -------
         data, filtered by range_obj
         """
+        if isinstance(range_obj, pd.DatetimeIndex):
+            range_obj = DateRange(range_obj.min(), range_obj.max())
         if 'date' in data.index.names:
             return data[(data.index.get_level_values('date') < range_obj.start) | (data.index.get_level_values('date') > range_obj.end)]
         elif 'date' in data.columns:

tests/integration/chunkstore/test_chunkstore.py

@@ -120,6 +120,8 @@ def test_write_read_with_daterange(chunkstore_lib):
     read_df = chunkstore_lib.read('test_df', chunk_range=DateRange(dt(2016, 1, 1), dt(2016, 1, 2)))
 
     assert_frame_equal(read_df, dg)
+    read_with_dr = chunkstore_lib.read('test_df', chunk_range=pd.date_range(dt(2016, 1, 1), dt(2016, 1, 2)))
+    assert_frame_equal(read_df, dg)
 
 
 def test_write_read_with_daterange_noindex(chunkstore_lib):

tests/unit/chunkstore/test_date_chunker.py

@@ -3,6 +3,7 @@
 from datetime import datetime as dt
 from arctic.date import DateRange
 from pandas.util.testing import assert_frame_equal
+import pandas as pd
 
 
 def test_date_filter():
@@ -28,3 +29,16 @@ def test_date_filter():
     assert_frame_equal(c.filter(df, DateRange(None, dt(2020, 1, 1))), df)
     # CLOSED - CLOSED (after range)
     assert(c.filter(df, DateRange(dt(2017, 1, 1), dt(2018, 1, 1))).empty)
+
+
+def test_date_filter_with_pd_date_range():
+    c = DateChunker()
+    df = DataFrame(data={'data': [1, 2, 3]},
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 1, 3), 1)],
+                                                names=['date', 'id'])
+                   )
+
+    assert(c.filter(df, pd.date_range(dt(2017, 1, 1), dt(2018, 1, 1))).empty)
+    assert_frame_equal(c.filter(df, pd.date_range(dt(2016, 1, 1), dt(2017, 1, 1))), df)