Fixes pandas-dev#171 and pandas-dev#172

Author: bmoscon

arctic/chunkstore/chunkstore.py

@@ -73,18 +73,32 @@ def __str__(self):
     def __repr__(self):
         return str(self)
 
-    def delete(self, symbol):
+    def delete(self, symbol, chunk_range=None):
         """
-        Delete all chunks for a symbol.
+        Delete all chunks for a symbol, or optionally, chunks within a range
 
         Parameters
         ----------
         symbol : `str`
             symbol name for the item
         """
-        query = {"symbol": symbol}
-        self._collection.delete_many(query)
-        self._collection.symbols.delete_many(query)
+        if chunk_range:
+            # read out chunks that fall within the range and filter out
+            # data within the range
+            df = self.read(symbol, chunk_range=chunk_range, no_filter=True)
+            df = df[(df.index.get_level_values('date') < chunk_range[0]) | (df.index.get_level_values('date') > chunk_range[1])]
+
+
+            # remove chunks, and update any remaining data
+            query = {'symbol': symbol}
+            query.update(self.chunker.to_mongo(chunk_range))
+            self._collection.delete_many(query)
+            self.update(symbol, df)
+
+        else:
+            query = {"symbol": symbol}
+            self._collection.delete_many(query)
+            self._collection.symbols.delete_many(query)
 
     def list_symbols(self):
         """
@@ -99,7 +113,7 @@ def list_symbols(self):
     def _get_symbol_info(self, symbol):
         return self._symbols.find_one({'symbol': symbol})
 
-    def read(self, symbol, chunk_range=None):
+    def read(self, symbol, chunk_range=None, no_filter=False):
         """
         Reads data for a given symbol from the database.
 
@@ -110,6 +124,9 @@ def read(self, symbol, chunk_range=None):
         chunk_range: object
             corresponding range object for the specified chunker (for
             DateChunker it is a DateRange object)
+        no_filter: boolean
+            perform chunk level filtering on the data (see filter() in _chunker)
+            only applicable when chunk_range is specified
 
         Returns
         -------
@@ -124,7 +141,7 @@ def read(self, symbol, chunk_range=None):
                 }
 
         if chunk_range:
-            spec['start'] = self.chunker.to_mongo(chunk_range)
+            spec.update(self.chunker.to_mongo(chunk_range))
 
         segments = []
         for _, x in enumerate(self._collection.find(spec, sort=[('start', pymongo.ASCENDING)],)):
@@ -137,7 +154,7 @@ def read(self, symbol, chunk_range=None):
 
         data = deserialize(records, sym['type'])
 
-        if chunk_range is None:
+        if no_filter or chunk_range is None:
             return data
         return self.chunker.filter(data, chunk_range)
 

arctic/chunkstore/date_chunker.py

@@ -74,7 +74,15 @@ def to_start_end(self, data):
         return self._get_date_range(data)
 
     def to_mongo(self, range_obj):
-        return range_obj.mongo_query()
+        if range_obj[0] and range_obj[1]:
+            return {'$and': [{'start': {'$lte': range_obj[1]}}, {'end': {'$gte': range_obj[0]}}]}
+        elif range_obj[0]:
+            return {'end': {'$gte': range_obj[0]}}
+        elif range_obj[1]:
+            return {'start': {'$lte': range_obj[1]}}
+        else:
+            return {}
+
 
     def filter(self, data, range_obj):
         return data.ix[range_obj[0]:range_obj[1]]

tests/integration/chunkstore/test_chunkstore.py

@@ -135,7 +135,7 @@ def test_monthly_df(chunkstore_lib):
     chunkstore_lib.write('chunkstore_test', df, chunk_size='M')
     ret = chunkstore_lib.read('chunkstore_test', chunk_range=DateRange(dt(2016, 1, 1), dt(2016, 1, 2)))
     assert len(ret) == 2
-    # assert_frame_equal(df, chunkstore_lib.read('chunkstore_test') )
+    assert_frame_equal(df, chunkstore_lib.read('chunkstore_test'))
 
 
 def test_yearly_df(chunkstore_lib):
@@ -638,3 +638,51 @@ def test_dtype_mismatch_error(chunkstore_lib):
     assert('Dtype mismatch' in  str(e))
 
     assert_series_equal(s, chunkstore_lib.read('test'))
+
+
+def test_delete_range(chunkstore_lib):
+    df = DataFrame(data={'data': [1, 2, 3, 4, 5, 6]},
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 2, 1), 1),
+                                                 (dt(2016, 2, 2), 1),
+                                                 (dt(2016, 3, 1), 1),
+                                                 (dt(2016, 3, 2), 1)],
+                                                names=['date', 'id'])
+                   )
+
+    df_result = DataFrame(data={'data': [1, 6]},
+                          index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                        (dt(2016, 3, 2), 1)],
+                                                       names=['date', 'id'])
+                          )
+
+    chunkstore_lib.write('test', df, 'M')
+    chunkstore_lib.delete('test', chunk_range=DateRange(dt(2016, 1, 2), dt(2016, 3, 1)))
+    assert_frame_equal(chunkstore_lib.read('test'), df_result)
+
+
+def test_read_chunk_range(chunkstore_lib):
+    df = DataFrame(data={'data': [1, 2, 3, 4, 5, 6, 7, 8, 9]},
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 1, 3), 1),
+                                                 (dt(2016, 2, 1), 1),
+                                                 (dt(2016, 2, 2), 1),
+                                                 (dt(2016, 2, 3), 1),
+                                                 (dt(2016, 3, 1), 1),
+                                                 (dt(2016, 3, 2), 1),
+                                                 (dt(2016, 3, 3), 1)],
+                                                names=['date', 'id'])
+                   )
+
+    chunkstore_lib.write('test', df, 'M')
+    assert(chunkstore_lib.read('test', chunk_range=DateRange(dt(2016, 1, 1), dt(2016, 1, 1))).index.get_level_values('date')[0] == dt(2016,1,1))
+    assert(chunkstore_lib.read('test', chunk_range=DateRange(dt(2016, 1, 2), dt(2016, 1, 2))).index.get_level_values('date')[0] == dt(2016, 1, 2))
+    assert(chunkstore_lib.read('test', chunk_range=DateRange(dt(2016, 1, 3), dt(2016, 1, 3))).index.get_level_values('date')[0] == dt(2016, 1, 3))
+    assert(chunkstore_lib.read('test', chunk_range=DateRange(dt(2016, 2, 2), dt(2016, 2, 2))).index.get_level_values('date')[0] == dt(2016, 2, 2))
+
+    assert(len(chunkstore_lib.read('test', chunk_range=DateRange(dt(2016, 2, 2), dt(2016, 2, 2)), no_filter=True)) == 3)
+
+    df2 = chunkstore_lib.read('test', chunk_range=DateRange(None, None))
+    assert_frame_equal(df, df2)